feat: add SessionWithParams interface for URL query parameter access #504
Conversation
## Walkthrough
A new `SessionWithParams` interface was introduced to allow session objects to expose a map of string parameters. The SSE and Streamable HTTP session implementations were updated to extract URL query parameters from incoming requests and store them in this map, accessible via a new `Params()` method. Comprehensive tests were added to verify correct parameter parsing, thread safety, and integration with session contexts.
## Changes
| Files/Paths | Change Summary |
|----------------------------------------------|-----------------------------------------------------------------------------------------------------------------|
| server/session.go | Added `SessionWithParams` interface extending `ClientSession` with `Params()` method. |
| server/session_test.go | Added `sessionTestClientWithParams` type, parameter-related methods, interface compliance check, and integration test for parameter handling and concurrency. |
| server/sse.go | Added `params` field and `Params()` method to `sseSession`; extracted query parameters in `handleSSE`. |
| server/sse_test.go | Added subtest to `TestSSEServer` for verifying parameter parsing from URL query and tool access. |
| server/streamable_http.go | Added `params` field and `Params()` method to `streamableHttpSession`; updated session creation to include query parameters. |
| server/streamable_http_test.go | Added `TestStreamableHTTP_ParameterParsing` to verify correct propagation and access of query parameters. |
## Estimated code review effort
3 (~45 minutes)
## Possibly related PRs
- mark3labs/mcp-go#273: Implements StreamableHTTPServer and streamableHttpSession; this PR extends it by adding parameter handling and the `Params()` method.
## Suggested labels
`type: enhancement`, `area: sdk`
## Suggested reviewers
- ezynda3
- robert-jackson-glean
- pottekkat📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (6)
🧠 Learnings (2)server/sse_test.go (7)Learnt from: robert-jackson-glean Learnt from: ezynda3 Learnt from: leavez Learnt from: xinwo Learnt from: octo Learnt from: xinwo Learnt from: floatingIce91 server/session_test.go (2)Learnt from: octo Learnt from: ezynda3 🚧 Files skipped from review as they are similar to previous changes (4)
🧰 Additional context used🧠 Learnings (2)server/sse_test.go (7)Learnt from: robert-jackson-glean Learnt from: ezynda3 Learnt from: leavez Learnt from: xinwo Learnt from: octo Learnt from: xinwo Learnt from: floatingIce91 server/session_test.go (2)Learnt from: octo Learnt from: ezynda3 🔇 Additional comments (4)
✨ Finishing Touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
server/session_test.go (1)
201-238: Thread-safe implementation with proper defensive copying.The
Params()andSetParams()methods correctly implement thread-safe access with defensive copying to prevent concurrent modification issues. The pattern matches other session types likeSessionWithTools.Minor: Consider removing redundant
GetParams()method.The
GetParams()method (lines 218-220) appears redundant since it just delegates toParams(). Unless this is specifically needed for testing purposes, it could be removed for simplicity.-// GetParams returns the current params (for testing purposes) -func (f *sessionTestClientWithParams) GetParams() map[string]string { - return f.Params() -}
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
server/session.go(1 hunks)server/session_test.go(3 hunks)server/sse.go(3 hunks)server/sse_test.go(1 hunks)server/streamable_http.go(5 hunks)server/streamable_http_test.go(1 hunks)
🧠 Learnings (4)
server/streamable_http_test.go (7)
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: 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: 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: 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.
server/sse_test.go (7)
Learnt from: robert-jackson-glean
PR: #214
File: server/sse.go:0-0
Timestamp: 2025-04-28T00:14:49.263Z
Learning: The SSE server in mcp-go implements path sanitization within the WithDynamicBasePath function that ensures the dynamic base path starts with "/" and has no trailing "/" to prevent double slashes in URL construction.
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.
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 the mark3labs/mcp-go project should handle both InputSchema and RawInputSchema consistently between MarshalJSON and UnmarshalJSON methods, even though the tools response from MCP server typically doesn't contain rawInputSchema.
Learnt from: floatingIce91
PR: #401
File: server/server.go:1082-1092
Timestamp: 2025-06-23T11:10:42.948Z
Learning: In Go MCP server, ServerTool.Tool field is only used for tool listing and indexing, not for tool execution or middleware. During handleToolCall, only the Handler field is used, so dynamic tools don't need the Tool field populated.
server/sse.go (1)
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.
server/session_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.
🧬 Code Graph Analysis (1)
server/streamable_http.go (1)
mcp/types.go (2)
JSONRPCNotification(319-322)Params(163-163)
🧰 Additional context used
🧠 Learnings (4)
server/streamable_http_test.go (7)
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: 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: 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: 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.
server/sse_test.go (7)
Learnt from: robert-jackson-glean
PR: #214
File: server/sse.go:0-0
Timestamp: 2025-04-28T00:14:49.263Z
Learning: The SSE server in mcp-go implements path sanitization within the WithDynamicBasePath function that ensures the dynamic base path starts with "/" and has no trailing "/" to prevent double slashes in URL construction.
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.
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 the mark3labs/mcp-go project should handle both InputSchema and RawInputSchema consistently between MarshalJSON and UnmarshalJSON methods, even though the tools response from MCP server typically doesn't contain rawInputSchema.
Learnt from: floatingIce91
PR: #401
File: server/server.go:1082-1092
Timestamp: 2025-06-23T11:10:42.948Z
Learning: In Go MCP server, ServerTool.Tool field is only used for tool listing and indexing, not for tool execution or middleware. During handleToolCall, only the Handler field is used, so dynamic tools don't need the Tool field populated.
server/sse.go (1)
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.
server/session_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.
🧬 Code Graph Analysis (1)
server/streamable_http.go (1)
mcp/types.go (2)
JSONRPCNotification(319-322)Params(163-163)
🔇 Additional comments (15)
server/session.go (1)
51-56: LGTM! Well-designed interface following existing patterns.The
SessionWithParamsinterface is cleanly designed, follows Go conventions, and is consistent with other session extension interfaces in the codebase. The method signature appropriately returns URL query parameters asmap[string]string.server/sse.go (3)
33-33: LGTM! Field addition is appropriate.The
paramsfield correctly stores URL query parameters for the session using the appropriatemap[string]stringtype.
52-54: LGTM! Clean implementation of the SessionWithParams interface.The
Params()method correctly exposes the session parameters as required by theSessionWithParamsinterface.
355-367: LGTM! Parameter extraction logic is well-implemented.The parameter extraction correctly captures URL query parameters during session creation. Note that when multiple values exist for the same parameter key, only the first value is retained (
v[0]), which is a reasonable design choice for simplicity.server/streamable_http_test.go (1)
899-1062: Excellent comprehensive test for parameter parsing functionality.This test thoroughly validates the
SessionWithParamsinterface implementation with excellent coverage:
- Good test design: Uses table-driven tests covering single parameters, multiple parameters, no parameters, and special characters
- Realistic testing approach: Tests through actual HTTP requests rather than unit testing in isolation
- Comprehensive validation: Checks both presence of expected parameters and absence of unexpected ones
- Edge case coverage: Includes URL encoding/decoding scenarios with special characters
- Proper structure: Clear setup, execution, and validation phases with appropriate cleanup
The test effectively ensures that URL query parameters are correctly parsed and accessible through the session context.
server/sse_test.go (1)
1610-1784: Comprehensive test for SessionWithParams interface.The test implementation is well-structured and covers multiple scenarios effectively. The test validates URL query parameter parsing and propagation through the SessionWithParams interface.
However, consider these potential improvements:
Thread Safety Concern: The global
capturedParamsvariable could cause issues if tests run concurrently. Consider using a channel or mutex for safer parameter capture.Multiple Query Parameter Values: The current implementation only captures the first value of each query parameter (
params[k] = v[0]). Consider testing scenarios with multiple values for the same key to ensure the behavior is documented and intentional.Here's a suggested improvement for safer parameter capture:
- var capturedParams map[string]string + capturedParams := make(chan map[string]string, 1) mcpServer.AddTool( mcp.NewTool("get-params"), func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { session := ClientSessionFromContext(ctx) if session == nil { return nil, fmt.Errorf("no session in context") } sessionWithParams, ok := session.(SessionWithParams) if !ok { return nil, fmt.Errorf("session does not implement SessionWithParams") } - capturedParams = sessionWithParams.Params() + select { + case capturedParams <- sessionWithParams.Params(): + default: + } return mcp.NewToolResultText("params captured"), nil }, )And update the verification section:
- // Verify captured params match expected - if capturedParams == nil { - t.Fatal("No params were captured") - } + var params map[string]string + select { + case params = <-capturedParams: + case <-time.After(time.Second): + t.Fatal("No params were captured within timeout") + }server/streamable_http.go (4)
240-243: Parameter extraction logic is correct but consider edge cases.The query parameter extraction correctly creates a map of string parameters from the URL query. However, consider these points:
Multiple Values: The current implementation only takes the first value (
v[0]) when a parameter has multiple values. This is a reasonable default, but it should be documented behavior.Empty Values: URL query parameters can have empty values (e.g.,
?key=), which will be stored as empty strings. This appears to be the intended behavior based on the test cases.The implementation aligns well with the test expectations and provides consistent behavior across different session types.
267-267: Correct parameter passing to session constructor.The session constructor call correctly includes the extracted query parameters, maintaining consistency with the updated constructor signature.
378-383: Consistent parameter extraction across request handlers.The parameter extraction logic in
handleGetmirrors the implementation inhandlePost, ensuring consistent behavior across both HTTP methods. The session constructor correctly receives the extracted parameters.This consistency is important for maintaining predictable behavior regardless of how clients interact with the server (POST requests vs. GET streaming).
600-600: Proper implementation of SessionWithParams interface.The changes correctly implement the SessionWithParams interface by:
- Adding a
paramsfield to store query parameters- Updating the constructor to accept and store parameters
- Providing a
Params()method for read-only accessConsider defensive copying for better encapsulation:
While the current implementation works correctly, consider making a defensive copy of the params map to prevent potential external modifications:
func newStreamableHttpSession(sessionID string, toolStore *sessionToolsStore, levels *sessionLogLevelsStore, params map[string]string) *streamableHttpSession { + // Make a defensive copy to prevent external modifications + paramsCopy := make(map[string]string, len(params)) + for k, v := range params { + paramsCopy[k] = v + } + s := &streamableHttpSession{ sessionID: sessionID, notificationChannel: make(chan mcp.JSONRPCNotification, 100), tools: toolStore, logLevels: levels, - params: params, + params: paramsCopy, } return s }However, since the current callers create new maps that aren't reused, this optimization may not be necessary unless external code might modify the maps after session creation.
Also applies to: 603-603, 609-609, 618-620
server/session_test.go (5)
176-199: LGTM! Consistent implementation pattern.The
sessionTestClientWithParamsstruct and basicClientSessioninterface methods follow the same well-established pattern as other session test clients in this file.
246-246: LGTM! Proper interface compliance verification.The addition of
SessionWithParamsinterface compliance check follows the established pattern and ensures compile-time verification.
1577-1622: Comprehensive test setup and basic functionality verification.The test properly:
- Creates realistic session parameters matching the PR objectives (tenant_id, user_id, environment)
- Verifies context storage and retrieval
- Tests interface compliance with
SessionWithParams- Validates parameter access and values
- Tests defensive copying behavior by attempting to modify returned params
1623-1668: Excellent concurrent access testing.The concurrent test properly validates thread safety by:
- Running multiple reader and writer goroutines simultaneously
- Using error channels to collect race condition issues
- Validating both nil checks and expected parameter counts
- Testing the RWMutex implementation under realistic concurrent load
This should effectively catch any race condition issues in the implementation.
1670-1684: Good edge case coverage.The subtests properly validate boundary conditions:
- Nil params correctly return nil (maintaining the distinction from empty)
- Empty params return an empty map rather than nil
- Both scenarios are important for robust parameter handling
- Add SessionWithParams interface to session.go for accessing URL query parameters - Implement Params() method in SSE and StreamableHTTP sessions to parse query parameters - Add comprehensive tests for parameter parsing with concurrent access protection - Support special characters and empty parameter scenarios This enables server.WithToolFilter and other middleware to access session-specific parameters like tenant_id, user_id, or environment from URL query strings, allowing for fine-grained filtering and context-aware tool execution.
Description
This enables server.WithToolFilter and other middleware to access session-specific parameters like tenant_id, user_id, or environment from URL query strings, allowing for fine-grained filtering and context-aware tool execution.
Type of Change
Checklist
Summary by CodeRabbit
New Features
Bug Fixes
Tests