feat: implement MCP elicitation support (#413) #548
Conversation
* Add ElicitationRequest, ElicitationResult, and related types to mcp/types.go * Implement server-side RequestElicitation method with session support * Add client-side ElicitationHandler interface and request handling * Implement elicitation in stdio and in-process transports * Add comprehensive tests following sampling patterns * Create elicitation example demonstrating usage patterns * Use 'Elicitation' prefix for type names to maintain clarity
WalkthroughAdds bidirectional elicitation across MCP: protocol types, client ElicitationHandler and option, server RequestElicitation and ServerOption, transport/session plumbing (stdio & in-process), tests, and an example demonstrating accept/decline/cancel flows and schema-bearing requests. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Assessment against linked issues
Out-of-scope changes
Possibly related PRs
Suggested labels
Suggested reviewers
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. 📜 Recent review detailsConfiguration used: CodeRabbit UI 💡 Knowledge Base configuration:
You can enable these settings in your CodeRabbit configuration. 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type 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: 5
🔭 Outside diff range comments (6)
client/transport/inprocess.go (4)
57-66: Ensure sessionID is generated before creating/registering the sessionIf NewInProcessTransport() (without options) is used, sessionID remains empty. Creating a session with an empty ID can break RegisterSession/UnregisterSession lookups and GetSessionId. Generate an ID when missing and persist it.
Apply this diff:
func (c *InProcessTransport) Start(ctx context.Context) error { - // Create and register session if we have handlers - if c.samplingHandler != nil || c.elicitationHandler != nil { - c.session = server.NewInProcessSessionWithHandlers(c.sessionID, c.samplingHandler, c.elicitationHandler) + // Create and register session if we have handlers + if c.samplingHandler != nil || c.elicitationHandler != nil { + // Ensure we have a session ID + if c.sessionID == "" { + c.sessionID = c.server.GenerateInProcessSessionID() + } + c.session = server.NewInProcessSessionWithHandlers(c.sessionID, c.samplingHandler, c.elicitationHandler) if err := c.server.RegisterSession(ctx, c.session); err != nil { return fmt.Errorf("failed to register session: %w", err) } } return nil }
94-103: Preserve session context for notificationsNotifications should also carry the client session in the context to keep behavior consistent with requests (logging level, client info, etc.).
Apply this diff:
func (c *InProcessTransport) SendNotification(ctx context.Context, notification mcp.JSONRPCNotification) error { notificationBytes, err := json.Marshal(notification) if err != nil { return fmt.Errorf("failed to marshal notification: %w", err) } notificationBytes = append(notificationBytes, '\n') + // Add session to context if available + if c.session != nil { + ctx = c.server.WithContext(ctx, c.session) + } c.server.HandleMessage(ctx, notificationBytes) return nil }
111-116: Use actual session ID when unregisteringIf sessionID wasn’t set on construction, UnregisterSession with an empty ID will fail. Fall back to the session’s ID.
Apply this diff:
func (c *InProcessTransport) Close() error { if c.session != nil { - c.server.UnregisterSession(context.Background(), c.sessionID) + sid := c.sessionID + if sid == "" { + sid = c.session.SessionID() + } + c.server.UnregisterSession(context.Background(), sid) } return nil }
118-120: Fix: Return actual session ID from GetSessionId()Currently always returns an empty string.
Apply this diff:
func (c *InProcessTransport) GetSessionId() string { - return "" + if c.sessionID != "" { + return c.sessionID + } + if c.session != nil { + return c.session.SessionID() + } + return "" }server/stdio.go (2)
223-225: Same stdout write race in RequestSamplingSame issue as RequestElicitation: direct writes to the shared stdout without coordinating with the server’s write lock. Use the same writeWithGlobalLock helper to serialize all writes.
Apply this diff:
- if _, err := writer.Write(requestBytes); err != nil { - return nil, fmt.Errorf("failed to write sampling request: %w", err) - } + if _, err := writeWithGlobalLock(writer, requestBytes); err != nil { + return nil, fmt.Errorf("failed to write sampling request: %w", err) + }
695-701: Unify server write path with the same global write lock to avoid interleaving with session-originated requestswriteResponse currently uses a different mutex (s.writeMu) than session-originated requests. To guarantee atomic frames across all writers, route all writes through the same global lock.
Apply this diff:
- // Protect concurrent writes - s.writeMu.Lock() - defer s.writeMu.Unlock() - - // Write response followed by newline - if _, err := fmt.Fprintf(writer, "%s\n", responseBytes); err != nil { - return err - } + // Write response followed by newline under shared lock + if _, err := writeWithGlobalLock(writer, append(responseBytes, '\n')); err != nil { + return err + }Optionally, remove s.writeMu if no other code depends on it after this change.
🧹 Nitpick comments (17)
mcp/types.go (3)
59-62: Introduce elicitation method: looks correct and consistent with naming.Constant name and value are consistent with the rest of the method enum. Non-blocking nit: consider aligning the comment URL to the spec path once a canonical spec page exists.
834-841: Prefer a more precise type for RequestedSchema to improve safety and interop.Using
anyworks, but you’ll lose type guarantees and may incur unnecessary reflection. Two better options:
- Use json.RawMessage to preserve arbitrary JSON without reflection.
- If the repo already standardizes on a JSON Schema type (e.g., the Tool InputSchema handling per our past learnings), reuse that type for consistency.
Suggested minimal change:
- RequestedSchema any `json:"requestedSchema"` + RequestedSchema json.RawMessage `json:"requestedSchema"`
842-860: Response envelope is well modeled; consider adding a helper validator.The tri-state response with optional Value is clear. Consider adding a small helper (outside this file) to validate that Value is present iff Type is “accept” and to enforce schema-validation errors vs user-decline/cancel semantics. This will make downstream handler code simpler and safer.
server/session.go (1)
55-61: SessionWithElicitation interface is clear; document cancel vs error semantics.Consider clarifying in a comment when to return an ElicitationResult with Type=cancel vs returning an error (e.g., transport interruption vs user-initiated cancel). This avoids ambiguity for implementers across transports.
client/elicitation.go (1)
9-19: ElicitationHandler contract looks good; add guidance on cancel/error handling.Interface is minimal and clear. Non-blocking: add a sentence that user-initiated cancellations should produce a result with Type=cancel, while transport or context cancellations should return an error. This aligns client behavior with server expectations.
server/elicitation.go (1)
10-25: Server-side elicitation request routing is soundCorrectly retrieves the client session from context and delegates only if SessionWithElicitation is implemented. Error messages match tests.
Consider defining exported sentinel errors for "no active session" and "session does not support elicitation" to enable callers to handle them programmatically without string matching in the future. Keep current messages for test compatibility.
client/client.go (1)
527-573: Transport-level elicitation handling is correct and mirrors sampling
- Validates handler presence
- Parses params safely via JSON round-trip
- Constructs mcp.ElicitationRequest and returns properly marshaled result
Sampling and elicitation transport handlers share identical request-param parsing and response-marshaling patterns. Consider extracting a small generic helper to reduce duplication and error-prone maintenance (example skeleton):
// outside the changed ranges – illustrative func (c *Client) handleBidi[TParams any, TResult any]( ctx context.Context, req transport.JSONRPCRequest, build func(TParams) (any, error), // constructs the MCP request to call the handler with call func(context.Context, any) (TResult, error), // invokes the specific handler ) (*transport.JSONRPCResponse, error) { var params TParams if req.Params != nil { b, err := json.Marshal(req.Params) if err != nil { return nil, fmt.Errorf("failed to marshal params: %w", err) } if err := json.Unmarshal(b, ¶ms); err != nil { return nil, fmt.Errorf("failed to unmarshal params: %w", err) } } mcpReq, err := build(params) if err != nil { return nil, err } result, err := call(ctx, mcpReq) if err != nil { return nil, err } rb, err := json.Marshal(result) if err != nil { return nil, fmt.Errorf("failed to marshal result: %w", err) } return &transport.JSONRPCResponse{ JSONRPC: mcp.JSONRPC_VERSION, ID: req.ID, Result: json.RawMessage(rb), }, nil }This keeps both flows uniform and easier to extend (e.g., future bidirectional methods).
server/elicitation_test.go (1)
114-172: Happy-path test validates Accept flow and structured ValueAsserts type and content of response payload; representative example of structured elicitation.
You could add an additional sub-assert that Value conforms to the declared schema in the test (e.g., check both fields exist and are strings) to better document intended behavior.
server/stdio.go (3)
629-683: Tighten response discrimination to avoid false positivesIf a message has an id but neither result nor error, this method will parse, attempt to unmarshal nil, and treat it as an elicitation response error. Align this with handleSamplingResponse by ignoring such frames.
Apply this diff:
- id, err := response.ID.Int64() - if err != nil { - return false - } + id, err := response.ID.Int64() + if err != nil || (response.Result == nil && response.Error == nil) { + return false + }
423-446: Potential goroutine leak on context cancel in readNextLineWhen ctx is cancelled, this returns "", nil but leaves the spawned goroutine blocked on ReadString. This can leak goroutines until stdin closes. Consider refactoring to avoid a goroutine per read or to unblock the reader on cancel (e.g., by using an interruptible reader, piping through an io.Pipe on shutdown, or a single long-lived goroutine that you signal to stop).
I can draft a refactor using a single read goroutine feeding a channel that the loop selects on, with proper shutdown semantics.
327-333: Nit: comment contradicts behaviorComment says “Default to discarding logs” but logs are directed to os.Stderr.
Apply this diff to clarify:
- ), // Default to discarding logs + ),Or update the comment to match the behavior.
client/inprocess_elicitation_test.go (5)
23-33: Prefer map[string]any over map[string]interface{} for consistencyThroughout the repository, any is used. Keep tests consistent to avoid noise in diffs and improve readability.
Apply this diff:
- Value: map[string]interface{}{ + Value: map[string]any{ "response": "User provided data", "accepted": true, },
39-53: Minor: use map[string]any consistently in schema definitionsThe file mixes map[string]any and map[string]interface{}. Align on any.
Apply this diff:
- Properties: map[string]any{ - "action": map[string]any{ + Properties: map[string]any{ + "action": map[string]any{ "type": "string", "description": "Action to perform", }, },And within the elicitation request:
- RequestedSchema: map[string]interface{}{ + RequestedSchema: map[string]any{ "type": "object", - "properties": map[string]interface{}{ - "confirm": map[string]interface{}{ + "properties": map[string]any{ + "confirm": map[string]any{ "type": "boolean", "description": "Confirm the action", }, - "details": map[string]interface{}{ + "details": map[string]any{ "type": "string", "description": "Additional details", }, },
193-197: Prefer public option to set handler rather than mutating unexported fieldSetting client.elicitationHandler directly couples the test to internal details. Use the functional option to exercise public API.
Apply this diff:
- client := NewClient(inProcessTransport) - client.elicitationHandler = handler + client := NewClient(inProcessTransport, WithElicitationHandler(handler))
80-93: Coverage: add a failing elicitation path testThe tool path already returns a CallToolResult with IsError true when RequestElicitation fails. Add a test that injects a failing handler or forces RequestElicitation to error and assert the error content and IsError=true.
I can add a subtest that makes the handler return an error and asserts the tool returns a text error and IsError=true.
171-173: String stabilityTests compare exact strings. Consider asserting on the response type (Accept/Decline/Cancel) and presence of substrings to reduce brittleness to minor text changes.
I can adjust the assertion to check Contains("User accepted") instead of full equality if desired.
client/elicitation_test.go (1)
84-123: Strengthen assertions by checking ID echo and response payloadAfter unmarshalling, also assert that the response type matches the test case and that the Request/Response IDs are echoed back.
For example:
if result != nil { if result.ID != request.ID { t.Errorf("expected response ID %v, got %v", request.ID, result.ID) } var elicitationResult mcp.ElicitationResult if err := json.Unmarshal(result.Result, &elicitationResult); err != nil { t.Errorf("failed to unmarshal result: %v", err) } else if tt.expectedError == "" && elicitationResult.Response.Type == "" { t.Error("empty response type") } }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these settings in your CodeRabbit configuration.
📒 Files selected for processing (13)
client/client.go(5 hunks)client/elicitation.go(1 hunks)client/elicitation_test.go(1 hunks)client/inprocess_elicitation_test.go(1 hunks)client/transport/inprocess.go(3 hunks)examples/elicitation/main.go(1 hunks)mcp/types.go(4 hunks)server/elicitation.go(1 hunks)server/elicitation_test.go(1 hunks)server/inprocess_session.go(4 hunks)server/server.go(3 hunks)server/session.go(1 hunks)server/stdio.go(6 hunks)
🧰 Additional context used
🧠 Learnings (2)
📚 Learning: 2025-03-04T06:59:43.882Z
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.
Applied to files:
mcp/types.go
📚 Learning: 2025-06-30T07:13:17.052Z
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.
Applied to files:
server/server.go
🧬 Code Graph Analysis (11)
client/elicitation.go (1)
mcp/types.go (2)
ElicitationRequest(829-832)ElicitationResult(843-850)
server/session.go (1)
mcp/types.go (2)
ElicitationRequest(829-832)ElicitationResult(843-850)
client/elicitation_test.go (5)
mcp/types.go (18)
ElicitationResult(843-850)ElicitationRequest(829-832)ElicitationResponse(853-859)ElicitationResponseTypeAccept(866-866)ElicitationResponseTypeDecline(868-868)ElicitationResponseTypeCancel(870-870)JSONRPCRequest(329-334)NewRequestId(263-265)MethodElicitationCreate(61-61)Params(181-181)Result(250-254)InitializeResult(434-448)Implementation(506-509)ServerCapabilities(476-503)JSONRPCNotification(337-340)InitializeRequest(418-422)InitializeParams(424-430)ClientCapabilities(459-471)client/elicitation.go (1)
ElicitationHandler(11-19)server/inprocess_session.go (1)
ElicitationHandler(19-21)client/client.go (3)
Client(17-29)WithElicitationHandler(50-54)NewClient(71-81)client/transport/inprocess.go (1)
WithElicitationHandler(32-36)
server/elicitation.go (2)
mcp/types.go (2)
ElicitationRequest(829-832)ElicitationResult(843-850)server/session.go (2)
ClientSessionFromContext(82-87)SessionWithElicitation(56-60)
server/inprocess_session.go (4)
mcp/types.go (3)
ElicitationRequest(829-832)ElicitationResult(843-850)JSONRPCNotification(337-340)client/elicitation.go (1)
ElicitationHandler(11-19)server/session.go (4)
ClientSession(11-20)SessionWithLogging(23-29)SessionWithClientInfo(43-53)SessionWithElicitation(56-60)server/sampling.go (1)
SessionWithSampling(42-45)
client/client.go (5)
client/elicitation.go (1)
ElicitationHandler(11-19)server/inprocess_session.go (1)
ElicitationHandler(19-21)client/transport/inprocess.go (1)
WithElicitationHandler(32-36)mcp/types.go (10)
JSONRPCRequest(329-334)JSONRPCResponse(343-347)MethodSamplingCreateMessage(877-877)MethodElicitationCreate(61-61)ElicitationParams(835-840)Params(181-181)ElicitationRequest(829-832)Request(172-175)JSONRPC_VERSION(115-115)Result(250-254)client/transport/interface.go (2)
JSONRPCRequest(57-62)JSONRPCResponse(64-73)
server/elicitation_test.go (4)
mcp/types.go (9)
JSONRPCNotification(337-340)ElicitationResult(843-850)ElicitationRequest(829-832)Params(181-181)ElicitationParams(835-840)ElicitationResponse(853-859)ElicitationResponseTypeAccept(866-866)ElicitationResponseType(862-862)ElicitationResponseTypeDecline(868-868)server/server.go (2)
NewMCPServer(307-333)WithElicitation(293-297)mcp/utils.go (1)
ToBoolPtr(861-863)server/session.go (1)
ClientSession(11-20)
server/server.go (1)
mcp/utils.go (1)
ToBoolPtr(861-863)
client/transport/inprocess.go (1)
server/inprocess_session.go (4)
SamplingHandler(14-16)ElicitationHandler(19-21)InProcessSession(23-33)NewInProcessSessionWithHandlers(43-50)
client/inprocess_elicitation_test.go (7)
mcp/types.go (13)
ElicitationRequest(829-832)ElicitationResult(843-850)ElicitationResponse(853-859)ElicitationResponseTypeAccept(866-866)Params(181-181)ElicitationParams(835-840)Content(941-943)TextContent(947-954)TextContent(956-956)ElicitationResponseTypeDecline(868-868)ElicitationResponseTypeCancel(870-870)LATEST_PROTOCOL_VERSION(105-105)ClientCapabilities(459-471)server/server.go (2)
WithElicitation(293-297)MCPServer(144-169)mcp/tools.go (5)
Tool(557-572)Description(851-855)ToolInputSchema(612-617)Properties(1093-1097)Required(859-863)client/elicitation.go (1)
ElicitationHandler(11-19)server/inprocess_session.go (1)
ElicitationHandler(19-21)client/client.go (3)
Client(17-29)WithElicitationHandler(50-54)NewClient(71-81)client/transport/inprocess.go (2)
NewInProcessTransportWithOptions(44-55)WithElicitationHandler(32-36)
server/stdio.go (3)
mcp/types.go (8)
JSONRPCNotification(337-340)ElicitationResult(843-850)ElicitationRequest(829-832)Params(181-181)ElicitationParams(835-840)JSONRPC_VERSION(115-115)MethodElicitationCreate(61-61)Result(250-254)server/session.go (2)
ClientSession(11-20)SessionWithElicitation(56-60)server/sampling.go (1)
SessionWithSampling(42-45)
🔇 Additional comments (27)
mcp/types.go (4)
469-471: Client capability flag added appropriately.Using a pointer to an empty struct mirrors existing capability flags (e.g., Sampling) and keeps the wire format minimal. No issues.
501-503: Server capability flag added appropriately.Matches the client-side shape and is correctly nullable for opt-in behavior in initialize.
825-832: ElicitationRequest wrapper shape is fine and consistent with other request types.Named
Paramsfield is consistent with many existing request wrappers in this file (e.g., ReadResourceRequest). Good placement and JSON tags.
864-871: Enum values are reasonable and easy to consume.No issues with the response type constants.
server/server.go (3)
180-186: Capability plumbing extended with elicitation (opt-in).Adding
elicitation *boolaligns with the existing pattern (e.g., logging, sampling). Looks good.
639-642: Initialize advertises elicitation correctly when enabled.Consistent with other capabilities. No concerns.
292-298: Elicitation option: end-to-end implementation verifiedWithElicitation is OK — the server method, session interface + implementations, method constant, and tests are present.
- server/elicitation.go: func (s *MCPServer) RequestElicitation(...)
- server/session.go: RequestElicitation declared on the Session interface
- server/stdio.go, server/inprocess_session.go: concrete RequestElicitation implementations
- mcp/types.go: MethodElicitationCreate = "elicitation/create"
- Tests: server/elicitation_test.go, client/elicitation_test.go, client/inprocess_elicitation_test.go
examples/elicitation/main.go (1)
46-50: RequestElicitation usage looks good.The example demonstrates the intended request/response flow and handles all response types. Once the safety checks above are added, this is a solid showcase.
client/transport/inprocess.go (2)
14-18: Addition of elicitation handler field looks correctWires in elicitation support alongside sampling in the in-process transport.
32-36: New option for elicitation handler is consistent with samplingPublic option mirrors existing WithSamplingHandler pattern. No issues.
client/client.go (4)
28-29: Client now tracks an elicitation handlerField addition is minimal and consistent with sampling.
48-55: WithElicitationHandler option is correct and documentedMirrors the sampling option with appropriate capability advertisement note.
179-183: Initialize advertises elicitation capability when handler is setThis aligns with capability negotiation and the new bidirectional request type.
467-474: Incoming request switch extended to handle elicitationApproach is consistent with sampling; default branch remains defensive.
server/inprocess_session.go (5)
18-22: ElicitationHandler interface is well-scopedClear separation of concerns and symmetric with SamplingHandler.
31-33: Session now tracks elicitation handlerRead/write guarded by RWMutex in RequestElicitation; good thread-safety practice.
43-50: Constructor with both handlers improves clarityNewInProcessSessionWithHandlers simplifies wiring from transports/tests.
119-129: RequestElicitation correctly delegates to handlerUses read lock, nil-check, and returns handler output. Error message consistent with tests.
138-143: Interface compliance assertions are completeCatches regressions at compile time.
server/elicitation_test.go (2)
84-112: Unsupported-session test reads cleanly and validates exact errorUses WithElicitation option and a session without elicitation support. Solid coverage.
174-239: No action needed — fakeSession is defined in testsThe repository contains type
fakeSessionat server/server_test.go:1802 (found by rg). The original concern about an undefined type is incorrect — do not apply the suggested diff.
- server/server_test.go:1802 —
type fakeSession struct { ... }Likely an incorrect or invalid review comment.
server/stdio.go (3)
302-307: LGTM: safe writer swapUsing defer to unlock improves exception-safety during writer updates.
317-321: LGTM: initialized pending maps for sampling and elicitationPrevents nil map writes at runtime.
520-524: Ordering is fine; responses are intercepted before server handlingProcessing sampling and elicitation responses prior to routing to the server avoids misclassification and unnecessary work.
client/inprocess_elicitation_test.go (1)
115-123: Nice: end-to-end in-process setup validates elicitation plumbingCreating a real MCP server with a tool that triggers elicitation exercises the path realistically.
client/elicitation_test.go (2)
26-78: Good breadth of casesTable covers no-handler, accept/decline/cancel, and handler error propagation. Nicely scoped.
192-226: Mock transport shape is sufficientStart/Close/no-op notification handler are adequate for these tests.
| // RequestElicitation sends an elicitation request to the client and waits for the response. | ||
| func (s *stdioSession) RequestElicitation(ctx context.Context, request mcp.ElicitationRequest) (*mcp.ElicitationResult, error) { | ||
| s.mu.RLock() | ||
| writer := s.writer | ||
| s.mu.RUnlock() | ||
|
|
||
| if writer == nil { | ||
| return nil, fmt.Errorf("no writer available for sending requests") | ||
| } | ||
|
|
||
| // Generate a unique request ID | ||
| id := s.requestID.Add(1) | ||
|
|
||
| // Create a response channel for this request | ||
| responseChan := make(chan *elicitationResponse, 1) | ||
| s.pendingMu.Lock() | ||
| s.pendingElicitations[id] = responseChan | ||
| s.pendingMu.Unlock() | ||
|
|
||
| // Cleanup function to remove the pending request | ||
| cleanup := func() { | ||
| s.pendingMu.Lock() | ||
| delete(s.pendingElicitations, id) | ||
| s.pendingMu.Unlock() | ||
| } | ||
| defer cleanup() | ||
|
|
||
| // Create the JSON-RPC request | ||
| jsonRPCRequest := struct { | ||
| JSONRPC string `json:"jsonrpc"` | ||
| ID int64 `json:"id"` | ||
| Method string `json:"method"` | ||
| Params mcp.ElicitationParams `json:"params"` | ||
| }{ | ||
| JSONRPC: mcp.JSONRPC_VERSION, | ||
| ID: id, | ||
| Method: string(mcp.MethodElicitationCreate), | ||
| Params: request.Params, | ||
| } | ||
|
|
||
| // Marshal and send the request | ||
| requestBytes, err := json.Marshal(jsonRPCRequest) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to marshal elicitation request: %w", err) | ||
| } | ||
| requestBytes = append(requestBytes, '\n') | ||
|
|
||
| if _, err := writer.Write(requestBytes); err != nil { | ||
| return nil, fmt.Errorf("failed to write elicitation request: %w", err) | ||
| } | ||
|
|
||
| // Wait for the response or context cancellation | ||
| select { | ||
| case <-ctx.Done(): | ||
| return nil, ctx.Err() | ||
| case response := <-responseChan: | ||
| if response.err != nil { | ||
| return nil, response.err | ||
| } | ||
| return response.result, nil | ||
| } | ||
| } |
There was a problem hiding this comment.
Critical: concurrent writes to stdout are not synchronized in RequestElicitation, risking interleaved/broken JSON
RequestElicitation writes directly to the shared writer without coordinating with the server's write lock. Meanwhile, server responses and notifications use StdioServer.writeResponse with a different lock (s.writeMu). This can interleave bytes across concurrent writes (sampling requests, elicitation requests, server responses, notifications), corrupting frames.
Use a single shared lock for all writes to the underlying stdout writer.
Apply this diff to write under a shared lock:
- if _, err := writer.Write(requestBytes); err != nil {
- return nil, fmt.Errorf("failed to write elicitation request: %w", err)
- }
+ if _, err := writeWithGlobalLock(writer, requestBytes); err != nil {
+ return nil, fmt.Errorf("failed to write elicitation request: %w", err)
+ }Additionally, add a shared write helper and mutex (outside the selected range):
// At package level (near other vars)
var stdioWriteMu sync.Mutex
// writeWithGlobalLock ensures all writes to the stdio transport are serialized.
func writeWithGlobalLock(w io.Writer, b []byte) (int, error) {
stdioWriteMu.Lock()
defer stdioWriteMu.Unlock()
return w.Write(b)
}And update writeResponse and RequestSampling to use the same helper (see separate comments).
🤖 Prompt for AI Agents
In server/stdio.go around lines 239 to 300, RequestElicitation writes directly
to the shared writer without using the global stdio write lock, which can
interleave bytes with other writers; add a package-level mutex stdioWriteMu and
helper writeWithGlobalLock(w io.Writer, b []byte) that locks, writes and
unlocks, then replace the direct writer.Write(requestBytes) call with
writeWithGlobalLock(writer, requestBytes); also ensure writeResponse and
RequestSampling are updated to use the same helper so all stdout writes are
serialized.
There was a problem hiding this comment.
@JBUinfo You're right that s.mu.RLock() is being used, but it only protects access to the writer field itself, not the actual write operation to stdout.
The issue is that once RequestElicitation gets the writer reference (safely with s.mu.RLock()), it then calls writer.Write(requestBytes) without any coordination with other parts of the system that also write to stdout.
Meanwhile, StdioServer.writeResponse uses a separate s.writeMu lock for its writes. This means:
RequestElicitationwrites to stdout (unprotected write operation)writeResponsewrites to stdout (protected bys.writeMu)- These can happen concurrently, potentially interleaving bytes and corrupting JSON frames
The s.mu.RLock() protects the field access but not the I/O operation itself. We need all writes to stdout to use the same synchronization mechanism to prevent concurrent writes from corrupting the output stream.
ฅ(•ﻌ•)ฅ
✏️ Learnings added
Learnt from: JBUinfo
PR: mark3labs/mcp-go#548
File: server/stdio.go:239-300
Timestamp: 2025-08-14T16:35:34.070Z
Learning: In server/stdio.go, the s.mu mutex protects access to the writer field, but does not protect the actual write operations to stdout. Concurrent write operations to the same stdout stream need to be synchronized using the same mutex to prevent byte interleaving and JSON corruption.
Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
server/session.go (1)
55-61: Interface addition is appropriate; clarify cancellation semantics and fix doc grammar.Good place for the elicitation extension. Tighten the comment to specify context behavior and correct wording.
-// SessionWithElicitation is an extension of ClientSession that can send elicitation requests +// SessionWithElicitation is an extension of ClientSession that can send elicitation requests. type SessionWithElicitation interface { ClientSession - // RequestElicitation sends an elicitation request to the client and waits for response + // RequestElicitation sends an elicitation request to the client and waits for a response. + // Implementations MUST honor ctx cancellation/deadline and return an error consistent with ctx.Err(). RequestElicitation(ctx context.Context, request mcp.ElicitationRequest) (*mcp.ElicitationResult, error) }Optionally, add a typed accessor to reduce repetitive type assertions at call sites:
// SessionWithElicitationFromContext retrieves a SessionWithElicitation from context, if available. func SessionWithElicitationFromContext(ctx context.Context) (SessionWithElicitation, bool) { cs := ClientSessionFromContext(ctx) swe, ok := cs.(SessionWithElicitation) return swe, ok }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these settings in your CodeRabbit configuration.
📒 Files selected for processing (1)
server/session.go(1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
server/session.go (1)
mcp/types.go (2)
ElicitationRequest(829-832)ElicitationResult(843-850)
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (3)
examples/elicitation/main.go (3)
56-60: Good: guarded assertion prevents panics on malformed elicitation response.This addresses prior feedback about unsafe map assertions and avoids runtime panics.
121-131: Good: switched to safe parameter access.Using request.GetString("data", "") avoids unsafe type assertions and is robust when params are missing.
165-168: Don’t proceed on malformed accept payload; fail fast.Falling back to an empty map will treat malformed responses as declines/missing values. Return an error to make issues explicit.
- responseData, ok := result.Response.Value.(map[string]interface{}) - if !ok { - responseData = make(map[string]interface{}) - } + responseData, ok := result.Response.Value.(map[string]interface{}) + if !ok { + return nil, fmt.Errorf("unexpected response format") + }
🧹 Nitpick comments (4)
examples/elicitation/main.go (4)
22-42: Tighten schema: disallow extraneous fields.Add additionalProperties: false so clients don’t send unexpected fields.
RequestedSchema: map[string]interface{}{ "type": "object", "properties": map[string]interface{}{ "projectName": map[string]interface{}{ "type": "string", "description": "Name of the project", "minLength": 1, }, "framework": map[string]interface{}{ "type": "string", "description": "Frontend framework to use", "enum": []string{"react", "vue", "angular", "none"}, }, "includeTests": map[string]interface{}{ "type": "boolean", "description": "Include test setup", "default": true, }, }, + "additionalProperties": false, "required": []string{"projectName"}, },
157-163: Differentiate Decline vs Cancel; return an error on cancellation.Lumping all non-accept responses into “cancelled” loses semantics and forces clients to parse text. Treat Decline as a non-error result and Cancel as an error.
- if result.Response.Type != mcp.ElicitationResponseTypeAccept { - return &mcp.CallToolResult{ - Content: []mcp.Content{ - mcp.NewTextContent("Processing cancelled by user"), - }, - }, nil - } + switch result.Response.Type { + case mcp.ElicitationResponseTypeAccept: + // continue to validate fields below + case mcp.ElicitationResponseTypeDecline: + return &mcp.CallToolResult{ + Content: []mcp.Content{ + mcp.NewTextContent("Processing declined by user"), + }, + }, nil + case mcp.ElicitationResponseTypeCancel: + return nil, fmt.Errorf("processing cancelled by user") + default: + return nil, fmt.Errorf("unexpected response type: %s", result.Response.Type) + }
136-149: Tighten schema for confirmation request.Disallow extra fields to keep the contract tight.
RequestedSchema: map[string]interface{}{ "type": "object", "properties": map[string]interface{}{ "proceed": map[string]interface{}{ "type": "boolean", "description": "Whether to proceed with processing", }, "reason": map[string]interface{}{ "type": "string", "description": "Optional reason for your decision", }, }, + "additionalProperties": false, "required": []string{"proceed"}, },
202-204: Handle SIGTERM for graceful shutdown in more environments.Add syscall.SIGTERM to support container orchestrators.
- signal.Notify(sigChan, os.Interrupt) + signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)Additional change outside the selected range:
// in the imports block "syscall"
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these settings in your CodeRabbit configuration.
📒 Files selected for processing (2)
client/elicitation_test.go(1 hunks)examples/elicitation/main.go(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- client/elicitation_test.go
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-06-20T20:39:51.870Z
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.
Applied to files:
examples/elicitation/main.go
🧬 Code Graph Analysis (1)
examples/elicitation/main.go (5)
server/server.go (4)
MCPServer(144-169)ToolHandlerFunc(41-41)NewMCPServer(307-333)WithElicitation(293-297)mcp/tools.go (2)
CallToolRequest(54-58)CallToolResult(40-51)mcp/types.go (7)
ElicitationRequest(829-832)Params(181-181)ElicitationParams(835-840)ElicitationResponseTypeAccept(866-866)Content(941-943)ElicitationResponseTypeDecline(868-868)ElicitationResponseTypeCancel(870-870)mcp/utils.go (1)
NewTextContent(199-204)server/stdio.go (1)
NewStdioServer(325-336)
| if proceed, ok := responseData["proceed"].(bool); !ok || !proceed { | ||
| reason := "No reason provided" | ||
| if r, ok := responseData["reason"].(string); ok && r != "" { | ||
| reason = r | ||
| } | ||
| return &mcp.CallToolResult{ | ||
| Content: []mcp.Content{ | ||
| mcp.NewTextContent(fmt.Sprintf("Processing declined: %s", reason)), | ||
| }, | ||
| }, nil | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Validate required field ‘proceed’ on accept path.
Since the schema requires proceed, treat missing/invalid types as an error rather than silently declining.
- if proceed, ok := responseData["proceed"].(bool); !ok || !proceed {
+ proceed, ok := responseData["proceed"].(bool)
+ if !ok {
+ return nil, fmt.Errorf("invalid or missing 'proceed' in elicitation response")
+ }
+ if !proceed {
reason := "No reason provided"
if r, ok := responseData["reason"].(string); ok && r != "" {
reason = r
}
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.NewTextContent(fmt.Sprintf("Processing declined: %s", reason)),
},
}, nil
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if proceed, ok := responseData["proceed"].(bool); !ok || !proceed { | |
| reason := "No reason provided" | |
| if r, ok := responseData["reason"].(string); ok && r != "" { | |
| reason = r | |
| } | |
| return &mcp.CallToolResult{ | |
| Content: []mcp.Content{ | |
| mcp.NewTextContent(fmt.Sprintf("Processing declined: %s", reason)), | |
| }, | |
| }, nil | |
| } | |
| proceed, ok := responseData["proceed"].(bool) | |
| if !ok { | |
| return nil, fmt.Errorf("invalid or missing 'proceed' in elicitation response") | |
| } | |
| if !proceed { | |
| reason := "No reason provided" | |
| if r, ok := responseData["reason"].(string); ok && r != "" { | |
| reason = r | |
| } | |
| return &mcp.CallToolResult{ | |
| Content: []mcp.Content{ | |
| mcp.NewTextContent(fmt.Sprintf("Processing declined: %s", reason)), | |
| }, | |
| }, nil | |
| } |
🤖 Prompt for AI Agents
In examples/elicitation/main.go around lines 170 to 180, the code currently
treats a missing or non-boolean "proceed" value as a silent decline; because the
schema requires "proceed", validate that responseData contains a boolean
"proceed" and if it's missing or the type is invalid return an explicit error
(e.g. fmt.Errorf with a clear message) instead of building a CallToolResult
decline; if "proceed" is present and false you may still return the existing
declined CallToolResult, but ensure missing/invalid values short-circuit with an
error response so callers can detect schema violations.
JBUinfo
force-pushed
the
simistrict-elicitation
branch
from
46e18dd to
2841d97
Compare
|
Hi @ezynda3 ! |
Exactly this PR #495 but fixing conflicts
Fixes #413
Summary by CodeRabbit
New Features
Documentation
Tests