Embed client capabilities into the Session #491
Conversation
Client sends its capabilities during the initialization step. This commit embeds the client capabilities into the client session in this step to enable subsequent executions are able to check the client capabilities to determine what actions they can perform. For instance, MCP Server checks the support of elicitation. If elicititation is supported by the client MCP Server can send elicitation request.
|
""" WalkthroughThis change introduces atomic storage and retrieval of client capabilities in session implementations. It extends session structs with a new atomic field and corresponding getter/setter methods, updates the session interface, and ensures client capabilities are set during initialization. Associated tests are updated to verify correct storage and retrieval of client capabilities. Additionally, documentation is added demonstrating client capability based filtering in an MCP server example. Changes
Estimated code review effort2 (~20 minutes) Possibly related PRs
Suggested labels
Suggested reviewers
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
server/sse.go (1)
33-33: Fix typo in comment.There's a typo in the comment: "capaabilities" should be "capabilities".
- clientCapabilities atomic.Value // stores session-specific client capaabilities + clientCapabilities atomic.Value // stores session-specific client capabilitiesserver/stdio.go (1)
104-115: Fix variable naming for consistency.The implementation follows established patterns correctly, but there's a variable naming issue in the getter method that should be fixed for clarity.
func (s *stdioSession) GetClientCapabilities() mcp.ClientCapabilities { if value := s.clientCapabilities.Load(); value != nil { - if clientInfo, ok := value.(mcp.ClientCapabilities); ok { - return clientInfo + if clientCapabilities, ok := value.(mcp.ClientCapabilities); ok { + return clientCapabilities } } return mcp.ClientCapabilities{} }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
server/inprocess_session.go(2 hunks)server/server.go(1 hunks)server/session.go(1 hunks)server/session_test.go(5 hunks)server/sse.go(2 hunks)server/stdio.go(2 hunks)
🧰 Additional context used
🧠 Learnings (7)
📓 Common learnings
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.
server/server.go (1)
Learnt from: ezynda3
PR: mark3labs/mcp-go#461
File: server/sampling.go:22-26
Timestamp: 2025-06-30T07:13:17.052Z
Learning: In the mark3labs/mcp-go project, the MCPServer.capabilities field is a struct value (serverCapabilities), not a pointer, so it cannot be nil and doesn't require nil checking. Only pointer fields within the capabilities struct should be checked for nil.
server/stdio.go (1)
Learnt from: ezynda3
PR: mark3labs/mcp-go#461
File: server/sampling.go:22-26
Timestamp: 2025-06-30T07:13:17.052Z
Learning: In the mark3labs/mcp-go project, the MCPServer.capabilities field is a struct value (serverCapabilities), not a pointer, so it cannot be nil and doesn't require nil checking. Only pointer fields within the capabilities struct should be checked for nil.
server/session.go (1)
Learnt from: ezynda3
PR: mark3labs/mcp-go#461
File: server/sampling.go:22-26
Timestamp: 2025-06-30T07:13:17.052Z
Learning: In the mark3labs/mcp-go project, the MCPServer.capabilities field is a struct value (serverCapabilities), not a pointer, so it cannot be nil and doesn't require nil checking. Only pointer fields within the capabilities struct should be checked for nil.
server/session_test.go (4)
Learnt from: ezynda3
PR: mark3labs/mcp-go#461
File: server/sampling.go:22-26
Timestamp: 2025-06-30T07:13:17.052Z
Learning: In the mark3labs/mcp-go project, the MCPServer.capabilities field is a struct value (serverCapabilities), not a pointer, so it cannot be nil and doesn't require nil checking. Only pointer fields within the capabilities struct should be checked for nil.
Learnt from: floatingIce91
PR: mark3labs/mcp-go#401
File: server/server.go:1082-1092
Timestamp: 2025-06-23T11:10:42.948Z
Learning: In Go MCP server, ServerTool.Tool field is only used for tool listing and indexing, not for tool execution or middleware. During handleToolCall, only the Handler field is used, so dynamic tools don't need the Tool field populated.
Learnt from: xinwo
PR: mark3labs/mcp-go#35
File: mcp/tools.go:107-137
Timestamp: 2025-03-04T06:59:43.882Z
Learning: Tool responses from the MCP server shouldn't contain RawInputSchema, which is why the UnmarshalJSON method for the Tool struct is implemented to handle only the structured InputSchema format.
Learnt from: octo
PR: mark3labs/mcp-go#149
File: mcptest/mcptest.go:0-0
Timestamp: 2025-04-21T21:26:32.945Z
Learning: In the mcptest package, prefer returning errors from helper functions rather than calling t.Fatalf() directly, giving callers flexibility in how to handle errors.
server/sse.go (2)
Learnt from: ezynda3
PR: mark3labs/mcp-go#461
File: server/sampling.go:22-26
Timestamp: 2025-06-30T07:13:17.052Z
Learning: In the mark3labs/mcp-go project, the MCPServer.capabilities field is a struct value (serverCapabilities), not a pointer, so it cannot be nil and doesn't require nil checking. Only pointer fields within the capabilities struct should be checked for nil.
Learnt from: 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/inprocess_session.go (1)
Learnt from: ezynda3
PR: mark3labs/mcp-go#461
File: server/sampling.go:22-26
Timestamp: 2025-06-30T07:13:17.052Z
Learning: In the mark3labs/mcp-go project, the MCPServer.capabilities field is a struct value (serverCapabilities), not a pointer, so it cannot be nil and doesn't require nil checking. Only pointer fields within the capabilities struct should be checked for nil.
🧬 Code Graph Analysis (5)
server/stdio.go (1)
mcp/types.go (1)
ClientCapabilities(441-451)
server/session.go (1)
mcp/types.go (1)
ClientCapabilities(441-451)
server/session_test.go (1)
mcp/types.go (3)
ClientCapabilities(441-451)InitializeRequest(400-404)Params(163-163)
server/sse.go (1)
mcp/types.go (1)
ClientCapabilities(441-451)
server/inprocess_session.go (1)
mcp/types.go (1)
ClientCapabilities(441-451)
🔇 Additional comments (11)
server/server.go (1)
586-586: LGTM! Client capabilities are properly stored during session initialization.The addition follows the established pattern for storing client info and is placed in the correct location within the initialization handler.
server/session.go (1)
49-53: LGTM! Interface extension follows established patterns.The new methods for client capabilities are consistent with the existing client info methods and provide a clean API for accessing capabilities across session types.
server/sse.go (1)
112-124: LGTM! Implementation follows established patterns.The getter and setter methods for client capabilities are implemented consistently with the existing client info methods, with proper thread safety and error handling.
server/inprocess_session.go (2)
24-24: LGTM! Field addition follows established patterns.The
clientCapabilitiesfield is properly added with atomic.Value for thread safety, consistent with other session implementations.
67-78: LGTM! Implementation follows established patterns.The getter and setter methods for client capabilities are implemented consistently with the existing client info methods, with proper thread safety and error handling.
server/stdio.go (1)
59-59: LGTM! Field addition follows established patterns.The
clientCapabilitiesfield is properly added with atomic.Value for thread safety, consistent with other session implementations.server/session_test.go (5)
109-109: LGTM: Consistent atomic field additionThe addition of the
clientCapabilities atomic.Valuefield follows the exact same pattern as the existingclientInfo atomic.Valuefield, maintaining consistency in the struct design.
141-152: LGTM: Well-implemented atomic getter and setter methodsThe implementation correctly follows the established pattern from the existing
GetClientInfo()andSetClientInfo()methods:
- Proper atomic operations for thread-safe access
- Safe type assertion with nil checks
- Appropriate default return value (empty struct) when no value is set
- Consistent error handling approach
905-905: LGTM: Minor formatting improvementGood formatting fix adding proper spacing in the test name.
1116-1123: LGTM: Appropriate test data initializationThe test correctly initializes
mcp.ClientCapabilitieswith a non-emptySamplingfield, which provides meaningful test data to verify the storage and retrieval functionality works correctly.
1146-1149: LGTM: Comprehensive test verificationThe test properly verifies that:
- The stored client capabilities can be retrieved via
GetClientCapabilities()- The retrieved capabilities match the original input data
- The atomic storage/retrieval mechanism works correctly
This ensures the new functionality integrates properly with the session initialization process.
|
Basically, after this PR, MCP Server will do something like this; session := ClientSessionFromContext(ctx)
if session == nil {
return nil, fmt.Errorf("no active session")
}
if clientSession, ok := session.(SessionWithClientInfo); ok {
clientSession.GetClientCapabilities()
} |
|
@ezynda3 would you mind taking a look at this, when you have a chance?. Thank you. |
|
Failure seems to be unrelated to the changes in this PR. |
|
Would you be able to add a bit of documentation to www/docs for this? |
|
I've updated the PR by adding documentation about this functionality. Please let me know what you think. Thank you. |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
www/docs/pages/servers/advanced.mdx (1)
837-841: Remove stray leading space befores.AddToolfor consistent formattingThe single extra space mis-aligns the call and breaks the otherwise uniform indentation in the snippet.
- s.AddTool( + s.AddTool(
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
www/docs/pages/servers/advanced.mdx(1 hunks)
🧠 Learnings (2)
📓 Common learnings
Learnt from: ezynda3
PR: mark3labs/mcp-go#461
File: server/sampling.go:22-26
Timestamp: 2025-06-30T07:13:17.052Z
Learning: In the mark3labs/mcp-go project, the MCPServer.capabilities field is a struct value (serverCapabilities), not a pointer, so it cannot be nil and doesn't require nil checking. Only pointer fields within the capabilities struct should be checked for nil.
www/docs/pages/servers/advanced.mdx (6)
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: 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: 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: 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: 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 context used
🧠 Learnings (2)
📓 Common learnings
Learnt from: ezynda3
PR: mark3labs/mcp-go#461
File: server/sampling.go:22-26
Timestamp: 2025-06-30T07:13:17.052Z
Learning: In the mark3labs/mcp-go project, the MCPServer.capabilities field is a struct value (serverCapabilities), not a pointer, so it cannot be nil and doesn't require nil checking. Only pointer fields within the capabilities struct should be checked for nil.
www/docs/pages/servers/advanced.mdx (6)
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: 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: 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: 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: 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.
www/docs/pages/servers/advanced.mdx
Outdated
| func handleCalculate(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { | ||
| session := server.ClientSessionFromContext(ctx) | ||
| if session == nil { | ||
| return nil, fmt.Errorf("no active session") | ||
| } | ||
|
|
||
| if clientSession, ok := session.(server.SessionWithClientInfo); ok { | ||
| clientCapabilities := clientSession.GetClientCapabilities() | ||
| if clientCapabilities.Sampling == nil { | ||
| fmt.Println("sampling is not enabled in client") | ||
| } | ||
| } | ||
| } | ||
| ``` |
There was a problem hiding this comment.
🛠️ Refactor suggestion
handleCalculate can fall through without a return → example will not compile
Go requires every control path in a non-void function to end with a return.
After the capability check the function falls through, triggering:
missing return at end of function
Provide a concrete result (or an explicit TODO) so users can copy-paste the snippet without errors.
func handleCalculate(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
@@
- if clientCapabilities.Sampling == nil {
- fmt.Println("sampling is not enabled in client")
- }
- }
-}
+ if clientCapabilities.Sampling == nil {
+ return mcp.NewToolResultError("sampling capability not enabled"), nil
+ }
+ }
+
+ // TODO: implement calculation logic
+ return mcp.NewToolResultError("not implemented"), 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.
| func handleCalculate(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { | |
| session := server.ClientSessionFromContext(ctx) | |
| if session == nil { | |
| return nil, fmt.Errorf("no active session") | |
| } | |
| if clientSession, ok := session.(server.SessionWithClientInfo); ok { | |
| clientCapabilities := clientSession.GetClientCapabilities() | |
| if clientCapabilities.Sampling == nil { | |
| fmt.Println("sampling is not enabled in client") | |
| } | |
| } | |
| } | |
| ``` | |
| func handleCalculate(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { | |
| session := server.ClientSessionFromContext(ctx) | |
| if session == nil { | |
| return nil, fmt.Errorf("no active session") | |
| } | |
| if clientSession, ok := session.(server.SessionWithClientInfo); ok { | |
| clientCapabilities := clientSession.GetClientCapabilities() | |
| if clientCapabilities.Sampling == nil { | |
| return mcp.NewToolResultError("sampling capability not enabled"), nil | |
| } | |
| } | |
| // TODO: implement calculation logic | |
| return mcp.NewToolResultError("not implemented"), nil | |
| } |
🤖 Prompt for AI Agents
In www/docs/pages/servers/advanced.mdx around lines 859 to 872, the function
handleCalculate lacks a return statement on all control paths, causing a
compilation error. Add a return statement at the end of the function that
returns a valid *mcp.CallToolResult and nil error, or include a TODO comment
with a placeholder return to ensure the function always returns a value.
Description
Client sends its capabilities during the initialization step.
This PR embeds the client capabilities into the client session in this step to enable subsequent executions are able to check the client capabilities to determine what actions they can perform. For instance, MCP Server checks the support of elicitation. If elicitation is supported by the client, MCP Server can send elicitation request.
Fixes ##144
Type of Change
Checklist
Summary by CodeRabbit