feat: implement sampling support for Streamable HTTP transport #515
Conversation
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 <noreply@anthropic.com>
WalkthroughThis change introduces full support for server-initiated sampling requests in the StreamableHTTP client transport, including bidirectional JSON-RPC request handling, error management, and response dispatch. It adds a new request handler mechanism, comprehensive unit tests for sampling flows, a working example client with a mock handler, and accompanying documentation. On the server side, it adds sampling support to the StreamableHTTPServer and sessions, enabling asynchronous sampling requests and responses over HTTP with SSE. The server capability declaration for sampling is added, and example server code and documentation are included. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Assessment against linked issues
Assessment against linked issues: Out-of-scope changes
Possibly related PRs
Suggested labels
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 (1)
🚧 Files skipped from review as they are similar to previous changes (1)
✨ 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. 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 (
|
andig
changed the title
|
Please be cautious: I've Claude do the work and only took a cursory glance. I cannot fully validate the result! |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (2)
examples/sampling_http_client/main.go (1)
104-107: Consider adding graceful shutdown handling.The current implementation blocks indefinitely. Consider adding signal handling for graceful shutdown in production use.
+import ( + "os" + "os/signal" + "syscall" +) // Keep the client running (in a real app, you'd have your main application logic here) +sigChan := make(chan os.Signal, 1) +signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) + select { case <-ctx.Done(): log.Println("Client context cancelled") +case <-sigChan: + log.Println("Received shutdown signal") }client/transport/streamable_http.go (1)
678-727: Consider context handling in async goroutine.The goroutine launched for handling requests uses the parent context, which might be cancelled before the request is fully processed. Consider creating a new context with timeout for the request handling.
// Handle the request in a goroutine to avoid blocking the SSE reader go func() { + // Create a new context with timeout for this specific request + reqCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + - response, err := handler(ctx, request) + response, err := handler(reqCtx, request) if err != nil { c.logger.Errorf("error handling request %s: %v", request.Method, err) // ... rest of error handling ... - c.sendResponseToServer(ctx, errorResponse) + c.sendResponseToServer(reqCtx, errorResponse) return } if response != nil { - c.sendResponseToServer(ctx, response) + c.sendResponseToServer(reqCtx, response) } }()
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
client/transport/streamable_http.go(5 hunks)client/transport/streamable_http_sampling_test.go(1 hunks)examples/sampling_http_client/README.md(1 hunks)examples/sampling_http_client/main.go(1 hunks)
🧰 Additional context used
🧠 Learnings (3)
📓 Common learnings
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/sampling_http_client/main.go (3)
Learnt from: ezynda3
PR: #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: davidleitw
PR: #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: #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/streamable_http_sampling_test.go (1)
Learnt from: octo
PR: #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 (1)
client/transport/streamable_http.go (3)
client/transport/interface.go (3)
RequestHandler(38-38)JSONRPCRequest(57-62)JSONRPCResponse(64-73)mcp/types.go (4)
Result(233-237)JSONRPCRequest(312-317)JSONRPCResponse(326-330)MethodSamplingCreateMessage(798-798)testdata/mockstdio_server.go (2)
JSONRPCRequest(13-18)JSONRPCResponse(20-28)
🔇 Additional comments (2)
examples/sampling_http_client/README.md (1)
1-95: Well-structured documentation!The README provides clear instructions, practical examples, and comprehensive coverage of the HTTP sampling client features. The separation between mock and real implementation guidance is particularly helpful.
examples/sampling_http_client/main.go (1)
59-62: Move defer statement after error check.The
deferstatement should be placed after the error check to avoid potential nil pointer dereference.if err != nil { log.Fatalf("Failed to create HTTP transport: %v", err) } -defer httpTransport.Close() +defer httpTransport.Close()⛔ Skipped due to learnings
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.
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 <noreply@anthropic.com>
andig
changed the title
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 <noreply@anthropic.com>
- 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 <noreply@anthropic.com>
This comment was marked as resolved.
This comment was marked as resolved.
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 5
♻️ Duplicate comments (2)
client/transport/streamable_http.go (2)
434-445: Request detection logic improved but consider additional validation.The updated logic using raw message parsing to check for the "method" field is more robust than the previous approach. However, consider adding validation for the JSON-RPC structure to ensure complete compliance.
The implementation correctly checks for the presence of the "method" field and non-nil ID to distinguish requests from responses, addressing the previous concern about relying solely on nil Result/Error fields.
731-758: Response sender includes required nil check.The sendResponseToServer method correctly includes nil checking for the response parameter as suggested in past reviews, preventing potential runtime panics.
The implementation properly handles HTTP POST for response delivery with appropriate error logging for non-OK status codes.
🧹 Nitpick comments (5)
server/sampling.go (1)
15-17: Consider simplifying the boolean assignment.The current implementation creates a local variable and takes its address. This can be simplified for better readability and to follow Go conventions.
-enabled := true -s.capabilities.sampling = &enabled +s.capabilities.sampling = mcp.ToBoolPtr(true)This uses the existing utility function
mcp.ToBoolPtrwhich is already used elsewhere in the codebase (like in the logging capability), providing consistency across the codebase.server/streamable_http_sampling_test.go (2)
16-44: Test should verify actual sampling functionality, not just initializationThis test only verifies that channels are initialized and the interface is implemented, but doesn't test the actual sampling request/response flow. Consider adding assertions that test sending a sampling request and receiving a response.
func TestStreamableHTTPServer_SamplingBasic(t *testing.T) { // Create MCP server with sampling enabled mcpServer := NewMCPServer("test-server", "1.0.0") mcpServer.EnableSampling() // Create HTTP server httpServer := NewStreamableHTTPServer(mcpServer) testServer := httptest.NewServer(httpServer) defer testServer.Close() // Test session creation and interface implementation sessionID := "test-session" session := newStreamableHttpSession(sessionID, httpServer.sessionTools, httpServer.sessionLogLevels) // Verify it implements SessionWithSampling _, ok := interface{}(session).(SessionWithSampling) if !ok { t.Error("streamableHttpSession should implement SessionWithSampling") } // Test that sampling request channels are initialized if session.samplingRequestChan == nil { t.Error("samplingRequestChan should be initialized") } if session.samplingResponseChan == nil { t.Error("samplingResponseChan should be initialized") } + + // Test actual sampling flow + ctx := context.Background() + request := mcp.CreateMessageRequest{ + CreateMessageParams: mcp.CreateMessageParams{ + Messages: []mcp.SamplingMessage{ + { + Role: mcp.RoleUser, + Content: mcp.TextContent{ + Type: "text", + Text: "Test sampling", + }, + }, + }, + }, + } + + // Simulate receiving a sampling response + go func() { + // Wait for the request to be sent + req := <-session.samplingRequestChan + // Send back a response + req.response <- samplingResponseItem{ + requestID: req.requestID, + result: &mcp.CreateMessageResult{ + Model: "test-model", + Role: mcp.RoleAssistant, + Content: mcp.TextContent{ + Type: "text", + Text: "Test response", + }, + }, + } + }() + + result, err := session.RequestSampling(ctx, request) + if err != nil { + t.Errorf("RequestSampling failed: %v", err) + } + if result == nil { + t.Error("Expected result, got nil") + } }
131-131: Avoid using nil for session stores in testsCreating a session with nil stores could cause panics if the implementation changes to use these stores. Consider using proper initialized stores for more robust tests.
- session := newStreamableHttpSession(sessionID, nil, nil) + session := newStreamableHttpSession(sessionID, newSessionToolsStore(), newSessionLogLevelsStore())examples/sampling_http_server/README.md (1)
49-62: Add language specifier to fenced code blockThe ASCII art diagram should have a language specifier for better syntax highlighting and to comply with markdown best practices.
-``` +```text Client (HTTP + SSE) ←→ Server (HTTP) │ │ ├─ POST: Tool Call ──→ │ │ │ │ ←── SSE: Sampling ───┤ │ Request │ │ │ ├─ POST: Sampling ───→ │ │ Response │ │ │ │ ←── HTTP: Tool ──────┤ Response</blockquote></details> <details> <summary>server/streamable_http.go (1)</summary><blockquote> `820-820`: **Improve error message for queue full condition** The error message could be more actionable for users. ```diff - return nil, fmt.Errorf("sampling request queue is full - server overloaded") + return nil, fmt.Errorf("sampling request queue is full (capacity: %d) - server overloaded, please retry later", cap(s.samplingRequestChan))
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
client/transport/streamable_http.go(5 hunks)client/transport/streamable_http_sampling_test.go(1 hunks)examples/sampling_http_server/README.md(1 hunks)examples/sampling_http_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/streamable_http.go(8 hunks)server/streamable_http_sampling_test.go(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- mcp/types.go
🚧 Files skipped from review as they are similar to previous changes (1)
- client/transport/streamable_http_sampling_test.go
🧰 Additional context used
🧠 Learnings (9)
📓 Common learnings
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.
server/sampling.go (1)
Learnt from: ezynda3
PR: #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.
server/server.go (1)
Learnt from: ezynda3
PR: #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.
server/sampling_test.go (2)
Learnt from: ezynda3
PR: #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: #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/streamable_http_sampling_test.go (2)
Learnt from: octo
PR: #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: #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/sampling_http_server/main.go (7)
Learnt from: floatingIce91
PR: #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: ezynda3
PR: #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: xinwo
PR: #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: #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: #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: octo
PR: #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: #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.
client/transport/streamable_http.go (4)
Learnt from: leavez
PR: #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: #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: xinwo
PR: #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: #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.
server/streamable_http.go (4)
Learnt from: xinwo
PR: #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: #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: ezynda3
PR: #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: leavez
PR: #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/sampling_http_server/README.md (2)
Learnt from: leavez
PR: #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: #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 (1)
client/transport/streamable_http.go (4)
client/transport/interface.go (3)
RequestHandler(38-38)JSONRPCRequest(57-62)JSONRPCResponse(64-73)mcp/types.go (3)
JSONRPCRequest(312-317)JSONRPCResponse(326-330)MethodSamplingCreateMessage(800-800)testdata/mockstdio_server.go (2)
JSONRPCRequest(13-18)JSONRPCResponse(20-28)client/transport/error.go (1)
Error(6-8)
🪛 markdownlint-cli2 (0.17.2)
examples/sampling_http_server/README.md
49-49: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
🔇 Additional comments (10)
client/transport/streamable_http.go (3)
112-114: LGTM: Thread-safe request handler support added.The addition of the request handler field with proper mutex protection enables safe concurrent access for server-to-client request handling.
567-572: LGTM: Proper request handler setter with thread safety.The SetRequestHandler method correctly implements thread-safe access to the request handler using the appropriate mutex.
654-729: Comprehensive request handling with proper error mapping.The handleIncomingRequest method implements solid asynchronous request processing with appropriate JSON-RPC error code mapping. The error handling correctly maps context errors to -32800 and other errors to -32603, following JSON-RPC 2.0 specifications.
The goroutine usage prevents blocking the SSE reader, which is crucial for maintaining the bidirectional communication flow.
server/server.go (2)
184-184: LGTM: Sampling capability field added correctly.The addition of the sampling boolean pointer field to serverCapabilities struct is consistent with the existing pattern used for other capabilities like logging.
584-586: LGTM: Proper conditional sampling capability advertisement.The initialization response correctly includes the sampling capability only when it's enabled, following the same pattern as other capabilities. The nil check ensures the capability is only advertised when explicitly enabled.
server/sampling_test.go (1)
117-154: LGTM: Comprehensive test for sampling capability enablement.The test properly verifies the EnableSampling functionality by:
- Confirming the capability is not set initially
- Enabling sampling via EnableSampling()
- Verifying the capability is then present in the initialization response
The test follows established patterns in the codebase and provides good coverage for the sampling capability feature.
examples/sampling_http_server/main.go (4)
15-18: LGTM: Proper server initialization with sampling enabled.The example correctly demonstrates server creation and sampling capability enablement following the established API patterns.
48-63: LGTM: Well-structured sampling request creation.The sampling request construction properly uses the MCP types and includes appropriate parameters like MaxTokens and Temperature for realistic LLM interaction.
65-81: LGTM: Proper timeout and error handling for sampling requests.The implementation correctly:
- Uses context with timeout for sampling requests (2 minutes is reasonable)
- Handles sampling errors gracefully by returning tool result with error flag
- Retrieves server from context using the proper API
This demonstrates good practices for production usage.
122-141: LGTM: Clear documentation and server setup.The example provides excellent documentation about:
- Server capabilities and requirements
- Client interaction patterns
- Available tools and their purposes
This helps users understand how to integrate with the sampling-enabled server.
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
- 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 <noreply@anthropic.com>
- 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 <noreply@anthropic.com>
- 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 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
server/streamable_http.go (2)
528-600: Implementation looks good but see concerns about delivery mechanismThe parsing and validation logic is well-implemented. However, the direct delivery to a session at line 591 is the area of concern mentioned by the PR author.
602-634: Acknowledge existing concern about response delivery reliabilityThe concerns about this implementation have been correctly identified in the previous review. The direct delivery mechanism is indeed fragile.
🧹 Nitpick comments (3)
server/streamable_http.go (3)
745-750: Unused fields in session structThe
mu sync.RWMutexandsamplingResponseChanappear to be unused in the implementation. Consider removing them if they're not needed.// Sampling support for bidirectional communication samplingRequestChan chan samplingRequestItem // server -> client sampling requests - samplingResponseChan chan samplingResponseItem // client -> server sampling responses samplingRequests sync.Map // requestID -> pending sampling request context requestIDCounter atomic.Int64 // for generating unique request IDs - mu sync.RWMutex // protects sampling channels and requests
755-760: Remove unused channel initializationFollowing up on the previous comment about unused fields.
sessionID: sessionID, notificationChannel: make(chan mcp.JSONRPCNotification, 100), tools: toolStore, logLevels: levels, samplingRequestChan: make(chan samplingRequestItem, 10), - samplingResponseChan: make(chan samplingResponseItem, 10),
812-851: Consider blocking send for sampling requestsThe implementation correctly uses per-request response channels. However, the non-blocking send at lines 832-839 could drop requests when the queue is full. Consider using a blocking send with context timeout instead.
// Send the sampling request via the channel (non-blocking) select { case s.samplingRequestChan <- samplingRequest: // Request queued successfully case <-ctx.Done(): return nil, ctx.Err() - default: - return nil, fmt.Errorf("sampling request queue is full - server overloaded") }This would make the send block until either the request is queued or the context is cancelled, preventing request loss during temporary congestion.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
server/streamable_http.go(8 hunks)server/streamable_http_sampling_test.go(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- server/streamable_http_sampling_test.go
🧰 Additional context used
🧠 Learnings (2)
📓 Common learnings
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.
server/streamable_http.go (4)
Learnt from: xinwo
PR: #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: #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: ezynda3
PR: #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: leavez
PR: #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.
🧬 Code Graph Analysis (1)
server/streamable_http.go (4)
mcp/types.go (11)
Result(233-237)MCPMethod(15-15)MethodInitialize(20-20)JSONRPCRequest(312-317)NewRequestId(246-248)Request(155-158)MethodSamplingCreateMessage(800-800)Params(164-164)CreateMessageParams(812-821)CreateMessageResult(827-834)CreateMessageRequest(807-810)client/transport/interface.go (1)
JSONRPCRequest(57-62)server/constants.go (1)
HeaderKeySessionID(5-5)server/sampling.go (1)
SessionWithSampling(42-45)
🔇 Additional comments (4)
server/streamable_http.go (4)
123-123: LGTM!The
sync.Mapis the appropriate choice for storing active sessions with concurrent access from multiple goroutines.
393-396: LGTM!Proper session lifecycle management with cleanup on disconnect.
425-439: LGTM!The sampling request is properly formatted as a JSON-RPC request and sent via the existing SSE infrastructure.
722-733: LGTM!Well-structured types for the sampling request/response flow.
- 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 <noreply@anthropic.com>
- 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 <noreply@anthropic.com>
- 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 <noreply@anthropic.com>
- 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 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
client/transport/streamable_http.go (1)
671-746: Comprehensive request handling with proper JSON-RPC compliance.This method excellently handles server-to-client requests with:
- Thread-safe handler access
- Appropriate JSON-RPC error responses for missing handlers
- Asynchronous processing to prevent SSE blocking
- Well-mapped error codes following JSON-RPC 2.0
Consider extracting error response creation to reduce duplication:
func createErrorResponse(id mcp.RequestId, code int, message string) *JSONRPCResponse { return &JSONRPCResponse{ JSONRPC: "2.0", ID: id, Error: &struct { Code int `json:"code"` Message string `json:"message"` Data json.RawMessage `json:"data"` }{ Code: code, Message: message, }, } }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
client/transport/streamable_http.go(7 hunks)examples/sampling_client/main.go(2 hunks)server/errors.go(1 hunks)server/streamable_http.go(8 hunks)
✅ Files skipped from review due to trivial changes (1)
- server/errors.go
🚧 Files skipped from review as they are similar to previous changes (1)
- server/streamable_http.go
🧰 Additional context used
🧠 Learnings (3)
📓 Common learnings
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/sampling_client/main.go (4)
Learnt from: octo
PR: #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: #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: #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: #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.
client/transport/streamable_http.go (4)
Learnt from: leavez
PR: #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: #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: xinwo
PR: #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: #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.
🧬 Code Graph Analysis (1)
client/transport/streamable_http.go (4)
client/transport/interface.go (3)
RequestHandler(38-38)JSONRPCRequest(57-62)JSONRPCResponse(64-73)mcp/types.go (3)
JSONRPCRequest(312-317)JSONRPCResponse(326-330)MethodSamplingCreateMessage(800-800)testdata/mockstdio_server.go (2)
JSONRPCRequest(13-18)JSONRPCResponse(20-28)client/transport/error.go (1)
Error(6-8)
🔇 Additional comments (9)
client/transport/streamable_http.go (6)
112-114: LGTM! Well-designed thread-safe request handling.The addition of
requestHandlerwith proper mutex protection enables safe concurrent access for server-to-client request handling. The use ofRWMutexis appropriate as multiple goroutines may need to read the handler while only one will set it.
403-408: Good defensive timeout handling.The conditional timeout application is well-designed - it respects existing shorter deadlines while ensuring requests don't hang indefinitely. The 30-second timeout aligns with the SSE timeout used elsewhere in the codebase.
441-452: Excellent improvement to request detection robustness.This implements the suggested improvement from the previous review. The new logic correctly identifies server requests by checking for the "method" field presence rather than relying on nil Result/Error fields, making it much more resilient to malformed messages.
574-579: LGTM! Consistent with existing patterns.The
SetRequestHandlermethod follows the same thread-safe pattern asSetNotificationHandler, using proper mutex locking for setting the handler.
598-625: Good improvements to connection reliability.The addition of connection timeouts and context-aware retry delays improves the robustness of the listening mechanism. The 10-second connection timeout prevents hanging on unresponsive servers, and the context-aware sleep enables immediate cancellation.
748-775: Excellent implementation with proper error handling.This method correctly implements the nil check suggested in the previous review and provides robust HTTP response handling. The context-aware request and comprehensive error logging ensure reliable response delivery.
examples/sampling_client/main.go (3)
8-9: LGTM! Standard imports for signal handling.The addition of
os/signalandsyscallpackages is appropriate for implementing graceful shutdown functionality.
95-106: Well-implemented graceful shutdown pattern.The signal handling follows Go best practices with a buffered channel, proper signal registration, and context-based cancellation propagation. The logging helps with observability during shutdown.
108-112: Good improvement to resource cleanup.Moving the defer after error checking ensures cleanup only occurs for successfully started clients. The error handling and logging for the Close operation provides better observability and debugging information.
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🔭 Outside diff range comments (1)
examples/sampling_client/main.go (1)
115-207: Update context usage for proper cancellation support.After implementing the graceful shutdown fix above, update all client operations to use the cancellable context (
opCtx) instead of the originalctx. This ensures that ongoing operations can be cancelled when a shutdown signal is received.Update the context parameter in these function calls:
mcpClient.Initialize(opCtx, ...)(line 115)mcpClient.ListTools(opCtx, ...)(line 135)mcpClient.CallTool(opCtx, ...)(lines 147, 168, 190)This will allow the client operations to be properly cancelled during graceful shutdown.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
client/transport/streamable_http.go(7 hunks)client/transport/streamable_http_sampling_test.go(1 hunks)examples/sampling_client/main.go(3 hunks)examples/sampling_server/main.go(1 hunks)server/session_test.go(3 hunks)server/streamable_http_sampling_test.go(1 hunks)server/streamable_http_test.go(2 hunks)
✅ Files skipped from review due to trivial changes (3)
- server/streamable_http_test.go
- examples/sampling_server/main.go
- server/session_test.go
🚧 Files skipped from review as they are similar to previous changes (3)
- server/streamable_http_sampling_test.go
- client/transport/streamable_http_sampling_test.go
- client/transport/streamable_http.go
🧰 Additional context used
🧠 Learnings (2)
📓 Common learnings
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/sampling_client/main.go (5)
Learnt from: octo
PR: #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: #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: #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: leavez
PR: #168
File: client/transport/streamable_http.go:183-186
Timestamp: 2025-04-17T15:58:33.813Z
Learning: In Go's atomic.Value, CompareAndSwap compares values using the == operator (value equality), not pointer identity. For strings, this means comparison by content, not memory address, so casting the loaded interface{} value to a string before comparison is unnecessary when using CompareAndSwap.
Learnt from: davidleitw
PR: #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.
🔇 Additional comments (2)
examples/sampling_client/main.go (2)
8-9: LGTM! Signal handling imports are appropriate.The addition of
os/signalandsyscallimports is necessary for implementing graceful shutdown functionality, which is a good practice for long-running client applications.
33-33: LGTM! Type alias update aligns with modern Go practices.The change from
map[string]interface{}tomap[string]anyis correct and follows Go 1.18+ conventions whereanyis the preferred alias forinterface{}.
andig
changed the title
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
server/streamable_http.go (1)
602-634: Implementation addresses key concerns but involves trade-offs.This method handles the core sampling response delivery with proper session lookup and request matching. The non-blocking delivery (
defaultcase) prevents deadlocks but trades off guaranteed delivery - responses may be lost if channels are full or sessions disconnect.For production use, consider implementing:
- Request timeout/cleanup mechanisms
- Response buffering for temporarily unavailable sessions
- Metrics for tracking dropped responses
The current approach prioritizes system stability over guaranteed delivery, which is a reasonable design choice for this transport layer.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
client/transport/streamable_http.go(7 hunks)examples/sampling_http_client/main.go(1 hunks)server/streamable_http.go(8 hunks)server/streamable_http_sampling_test.go(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- server/streamable_http_sampling_test.go
- examples/sampling_http_client/main.go
🧰 Additional context used
🧠 Learnings (2)
server/streamable_http.go (6)
Learnt from: xinwo
PR: #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: ezynda3
PR: #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: xinwo
PR: #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: octo
PR: #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: #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: leavez
PR: #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.
client/transport/streamable_http.go (4)
Learnt from: leavez
PR: #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: #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: xinwo
PR: #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: #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.
🧬 Code Graph Analysis (1)
client/transport/streamable_http.go (3)
client/transport/interface.go (3)
RequestHandler(38-38)JSONRPCRequest(57-62)JSONRPCResponse(64-73)mcp/types.go (3)
JSONRPCRequest(312-317)JSONRPCResponse(326-330)MethodSamplingCreateMessage(800-800)testdata/mockstdio_server.go (2)
JSONRPCRequest(13-18)JSONRPCResponse(20-28)
🔇 Additional comments (17)
client/transport/streamable_http.go (7)
112-114: LGTM! Good concurrency design.The addition of
requestHandlerfield with proper synchronization usingsync.RWMutexfollows established patterns in the codebase and enables safe concurrent access for server-to-client request handling.
403-408: LGTM! Consistent timeout handling.The 30-second timeout matches the established pattern from SSE client implementation and properly respects existing shorter deadlines in the context.
441-452: Excellent improvement in request detection robustness.This change addresses the previous review concern by using a more reliable method to detect server requests - checking for the presence of the "method" field rather than relying on Result/Error being nil. This makes the detection logic more robust against malformed messages.
574-579: LGTM! Consistent with existing patterns.The
SetRequestHandlermethod follows the same thread-safe pattern asSetNotificationHandlerand provides proper encapsulation for the request handler field.
598-625: Good timeout and cancellation improvements.The addition of connection timeouts and context-aware sleep improves the robustness of the listening loop by preventing indefinite blocking and enabling clean shutdown.
671-750: Well-implemented request handling with comprehensive error management.The
handleIncomingRequestmethod demonstrates good practices:
- Thread-safe handler access with read lock
- Asynchronous processing to avoid blocking the SSE reader
- Comprehensive JSON-RPC error code mapping for different scenarios
- Proper timeout handling with 30-second context
- Graceful error responses for missing handlers
The method properly handles the complexity of bidirectional JSON-RPC communication over HTTP.
752-779: Excellent implementation addressing previous review feedback.This method properly addresses the previous review concern by including a nil check for the response parameter. The error handling approach using logging is appropriate for this asynchronous operation, and the context management is correctly implemented.
server/streamable_http.go (10)
123-123: LGTM! Appropriate concurrent data structure.Using
sync.MapforactiveSessionsis the correct choice for tracking sessions with concurrent access from multiple goroutines handling sampling responses.
227-252: Solid message type detection logic.The JSON message parsing correctly differentiates between sampling responses (no method field) and regular requests (has method field). The detection logic is robust and handles the different message types appropriately.
393-396: Good session lifecycle management.Proper registration of sessions in
activeSessionswith cleanup via defer ensures that sessions are available for sampling response delivery and are properly cleaned up when the connection ends.
425-439: Well-integrated sampling request forwarding.The sampling request forwarding properly constructs JSON-RPC messages and integrates cleanly with the existing SSE mechanism while maintaining proper cancellation semantics.
528-600: Comprehensive sampling response handling.The
handleSamplingResponsemethod provides thorough validation of session IDs and request IDs, proper parsing of JSON-RPC responses (both success and error cases), and appropriate error handling throughout the flow.
722-733: Clean type definitions for sampling communication.The
samplingRequestItemandsamplingResponseItemstructs provide clear, well-structured data types for the bidirectional sampling communication flow with appropriate fields and types.
745-749: Well-designed session sampling state management.The addition of sampling-related fields uses appropriate concurrency primitives: atomic counter for request IDs,
sync.Mapfor request tracking, and proper mutex protection for the sampling channels and state.
753-759: Appropriate channel initialization.The constructor properly initializes the sampling request channel with a reasonable buffer size (10) to prevent blocking while maintaining bounded memory usage.
810-849: Excellent implementation of core sampling functionality.The
RequestSamplingmethod demonstrates good Go practices:
- Atomic request ID generation prevents collisions
- Proper resource cleanup with defer
- Non-blocking channel operations prevent deadlocks
- Comprehensive context cancellation handling
- Clean separation of concerns with dedicated response channels
This provides a solid foundation for server-initiated sampling requests.
851-851: Good interface compliance verification.Explicit declaration of
SessionWithSamplinginterface compliance ensures compile-time verification and makes the interface implementation clear to readers.
Removes unused sync.RWMutex field that was flagged by golangci-lint. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
Implements sampling capability for HTTP transport, resolving issue #419.
Todo
What's New
BidirectionalInterfacefor HTTP transportKey Features
Mcp-Session-Id,Mcp-Protocol-Version)Architecture
Testing
All tests pass ✅ including comprehensive coverage for:
Usage
Closes #419
Summary by CodeRabbit
Summary by CodeRabbit
New Features
Bug Fixes
Documentation
Tests
Chores