-
Notifications
You must be signed in to change notification settings - Fork 347
Manage tools on a per session basis #179
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
Caution Review failedThe pull request is closed. WalkthroughThis change refactors session management in the MCP server by moving all session-related interfaces, methods, and logic from Changes
Possibly related PRs
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (2)
✨ 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:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 6
♻️ Duplicate comments (1)
server/session.go (1)
213-214
: Duplicate unchecked error – same issue inDeleteSessionTools
The error from the notification helper is also ignored here. Please handle it as suggested
above to keep behaviour consistent and observable.🧰 Tools
🪛 golangci-lint (1.64.8)
214-214: Error return value of
s.SendNotificationToSpecificClient
is not checked(errcheck)
🧹 Nitpick comments (4)
README.md (2)
94-97
: Inconsistent bullet‑point style (markdown‑lint MD004)The table‑of‑contents list uses dashes (
-
) while the rest of the document (and the configured linter) expect asterisks (*
).
Although purely cosmetic, this currently triggers markdown‑lint violations in CI. Switching to a uniform list style will silence the warning.- - [Extras](#extras) - - [Session Management](#session-management) - - [Request Hooks](#request-hooks) - - [Tool Handler Middleware](#tool-handler-middleware) + * [Extras](#extras) + * [Session Management](#session-management) + * [Request Hooks](#request-hooks) + * [Tool Handler Middleware](#tool-handler-middleware)🧰 Tools
🪛 markdownlint-cli2 (0.17.2)
94-94: Unordered list style
Expected: asterisk; Actual: dash(MD004, ul-style)
95-95: Unordered list style
Expected: asterisk; Actual: dash(MD004, ul-style)
96-96: Unordered list style
Expected: asterisk; Actual: dash(MD004, ul-style)
97-97: Unordered list style
Expected: asterisk; Actual: dash(MD004, ul-style)
525-531
: Minor wording & compile‑time issues in code sampleThe snippet compiles for documentation only, but
server.WithToolCapabilities(true)
enables list‑changed notifications, not generic “session capabilities”.
Consider adding a clarifying comment or usingWithToolCapabilities(true)
+WithResourceCapabilities(true, true)
to avoid confusing newcomers.No code change required—just doc clarity.
🧰 Tools
🪛 markdownlint-cli2 (0.17.2)
526-526: Unordered list style
Expected: asterisk; Actual: dash(MD004, ul-style)
527-527: Unordered list style
Expected: asterisk; Actual: dash(MD004, ul-style)
528-528: Unordered list style
Expected: asterisk; Actual: dash(MD004, ul-style)
529-529: Unordered list style
Expected: asterisk; Actual: dash(MD004, ul-style)
server/session_test.go (1)
270-279
: Usestrings.HasPrefix
for clearer intentThe inline prefix check manually slices strings and can panic on short names (though you guarded with a length check). Using
strings.HasPrefix
is shorter and safer.- if len(tool.Name) >= len(prefix) && tool.Name[:len(prefix)] == prefix { + if strings.HasPrefix(tool.Name, prefix) {server/session.go (1)
86-95
: Notification drops are silently swallowedThe non‑blocking send prevents the server from being stalled by a slow client,
which is good. However, thedefault
branch currently discards the event
without visibility. Even a lightweight debug/trace log would help operators
spot problem clients.- default: - // TODO: log blocked channel in the future versions + default: + // channel full – consider logging or incrementing a metric
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
README.md
(2 hunks)server/errors.go
(1 hunks)server/server.go
(9 hunks)server/session.go
(1 hunks)server/session_test.go
(1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (3)
server/session_test.go (5)
mcp/types.go (4)
JSONRPCNotification
(192-195)Params
(104-104)Content
(663-665)Result
(173-177)server/server.go (3)
ServerTool
(50-53)NewMCPServer
(276-302)WithToolCapabilities
(252-259)server/session.go (3)
ClientSession
(10-19)SessionWithTools
(22-28)ClientSessionFromContext
(34-39)mcp/tools.go (7)
Tool
(69-80)NewTool
(144-166)CallToolRequest
(44-59)CallToolResult
(34-41)WithDescription
(187-191)ListToolsResult
(19-22)Description
(205-209)mcp/utils.go (1)
NewToolResultText
(214-223)
server/server.go (2)
mcp/tools.go (1)
Tool
(69-80)server/session.go (2)
ClientSessionFromContext
(34-39)SessionWithTools
(22-28)
server/session.go (4)
mcp/types.go (5)
JSONRPCNotification
(192-195)JSONRPC_VERSION
(81-81)Notification
(106-109)Params
(104-104)NotificationParams
(111-118)server/server.go (2)
ServerTool
(50-53)MCPServer
(126-151)server/errors.go (6)
ErrSessionExists
(16-16)ErrNotificationNotInitialized
(21-21)ErrNotificationChannelBlocked
(22-22)ErrSessionNotFound
(15-15)ErrSessionNotInitialized
(17-17)ErrSessionDoesNotSupportTools
(18-18)mcp/tools.go (1)
Tool
(69-80)
🪛 golangci-lint (1.64.8)
server/server.go
809-809: ineffectual assignment to ok
(ineffassign)
server/session.go
185-185: Error return value of s.SendNotificationToSpecificClient
is not checked
(errcheck)
214-214: Error return value of s.SendNotificationToSpecificClient
is not checked
(errcheck)
🪛 markdownlint-cli2 (0.17.2)
README.md
94-94: Unordered list style
Expected: asterisk; Actual: dash
(MD004, ul-style)
95-95: Unordered list style
Expected: asterisk; Actual: dash
(MD004, ul-style)
96-96: Unordered list style
Expected: asterisk; Actual: dash
(MD004, ul-style)
97-97: Unordered list style
Expected: asterisk; Actual: dash
(MD004, ul-style)
526-526: Unordered list style
Expected: asterisk; Actual: dash
(MD004, ul-style)
527-527: Unordered list style
Expected: asterisk; Actual: dash
(MD004, ul-style)
528-528: Unordered list style
Expected: asterisk; Actual: dash
(MD004, ul-style)
529-529: Unordered list style
Expected: asterisk; Actual: dash
(MD004, ul-style)
🔇 Additional comments (1)
server/errors.go (1)
7-23
: Centralized error declarations look solidThe sentinel errors are clearly grouped, exported, and ready to be wrapped/checked with
errors.Is
. No concerns from a correctness or maintainability standpoint.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 0
♻️ Duplicate comments (6)
server/session.go (3)
62-69
:⚠️ Potential issue
UnregisterSession
can panic when the session is unknown
LoadAndDelete
returns(nil, false)
when the key does not exist; the subsequent type‑assertion on anil
interface value panics.
This exact problem was raised in a previous review but is still present.-func (s *MCPServer) UnregisterSession( - ctx context.Context, - sessionID string, -) { - session, _ := s.sessions.LoadAndDelete(sessionID) - s.hooks.UnregisterSession(ctx, session.(ClientSession)) -} +func (s *MCPServer) UnregisterSession( + ctx context.Context, + sessionID string, +) { + sessionValue, ok := s.sessions.LoadAndDelete(sessionID) + if !ok { + return // nothing to unregister + } + if session, ok := sessionValue.(ClientSession); ok { + s.hooks.UnregisterSession(ctx, session) + } +}
161-181
: 🛠️ Refactor suggestionConcurrent write to the underlying
map[string]ServerTool
– possible data race
AddSessionTools
fetches the map viaGetSessionTools()
and mutates it in‑place without any synchronisation.
If two goroutines callAddSessionTools
/DeleteSessionTools
concurrently for the same session (or a test mutates the map directly), the map can be corrupted, leading to a panic.Consider:
- Guarding the per‑session map with a mutex inside the session implementation, or
- Cloning the map before mutation (
maps.Clone
in Go 1.21), then replacing it viaSetSessionTools
.At minimum, document the expectation that the caller provides external synchronisation.
184-186
:⚠️ Potential issueIgnored errors from
SendNotificationToSpecificClient
may hide delivery failuresBoth
AddSessionTools
andDeleteSessionTools
drop the error returned bySendNotificationToSpecificClient
.
If the target session is uninitialised or its channel is full, the notification is silently lost, contradicting the contract of these methods.- s.SendNotificationToSpecificClient(sessionID, "notifications/tools/list_changed", nil) + if err := s.SendNotificationToSpecificClient( + sessionID, + "notifications/tools/list_changed", + nil, + ); err != nil { + return err + }Propagate (or at least log) the error so callers can react appropriately.
Also applies to: 213-215
🧰 Tools
🪛 golangci-lint (1.64.8)
185-185: Error return value of
s.SendNotificationToSpecificClient
is not checked(errcheck)
server/server.go (2)
769-775
:⚠️ Potential issuePotential deadlock: filters executed while read‑lock is held
handleListTools
keepstoolFiltersMu.RLock()
while invoking every filter.
If a filter ever registers another filter (callsWithToolFilter
) or otherwise needs the write‑lock, the server deadlocks.This was pointed out earlier but not addressed. Copy the slice under the lock, then release it before executing user‑supplied code:
- s.toolFiltersMu.RLock() - if len(s.toolFilters) > 0 { - for _, filter := range s.toolFilters { - tools = filter(ctx, tools) - } - } - s.toolFiltersMu.RUnlock() + var filtersCopy []ToolFilterFunc + s.toolFiltersMu.RLock() + filtersCopy = append(filtersCopy, s.toolFilters...) + s.toolFiltersMu.RUnlock() + + for _, filter := range filtersCopy { + tools = filter(ctx, tools) + }
801-813
:⚠️ Potential issueInner
ok
shadows outerok
– session tools never winInside
handleToolCall
, the short variable declarationif sessionWithTools, ok := session.(SessionWithTools); ok {creates a new
ok
, leaving the outerok
unchanged (false
).
Consequently, even when a session‑specific tool exists, the function later falls back to the global map and may return “tool not found”.Fix by using a different name for the assertion result:
- if sessionWithTools, ok := session.(SessionWithTools); ok { + if sessionWithTools, okSession := session.(SessionWithTools); okSession {And update subsequent references (
ok = ...
) accordingly.
This also removes theineffassign
warning.🧰 Tools
🪛 golangci-lint (1.64.8)
809-809: ineffectual assignment to ok
(ineffassign)
server/session_test.go (1)
341-360
:⚠️ Potential issueTests won’t compile –
fakeSession
type is still missing
TestMCPServer_SendNotificationToSpecificClient
instantiates&fakeSession{…}
but no such type is defined in this package (onlysessionTestClient
andfakeSessionWithTools
exist).Either add the missing stub that satisfies
ClientSession
or reusesessionTestClient
:- session1 := &fakeSession{ + session1 := &sessionTestClient{Repeat for
session2
andsession3
.
Otherwise the entire test suite fails to build.
🧹 Nitpick comments (1)
server/session_test.go (1)
214-219
: Race‑prone flake: 100 ms timeout may be too aggressive on busy runnersThe notification checks use a 100 ms
time.After
. On slower CI machines this can produce false negatives.Consider increasing to ~500 ms or, better, use a context with deadline passed into the server so the test duration adapts to load.
Also applies to: 255-260
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
server/errors.go
(1 hunks)server/server.go
(8 hunks)server/session.go
(1 hunks)server/session_test.go
(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- server/errors.go
🧰 Additional context used
🧬 Code Graph Analysis (2)
server/session_test.go (4)
mcp/types.go (4)
JSONRPCNotification
(192-195)Params
(104-104)Content
(663-665)Result
(173-177)server/server.go (5)
ServerTool
(50-53)NewMCPServer
(276-302)WithToolCapabilities
(252-259)ToolFilterFunc
(47-47)WithToolFilter
(208-216)server/session.go (3)
ClientSession
(10-19)SessionWithTools
(22-28)ClientSessionFromContext
(34-39)mcp/tools.go (6)
Tool
(69-80)NewTool
(144-166)CallToolRequest
(44-59)CallToolResult
(34-41)WithDescription
(187-191)Description
(205-209)
server/server.go (2)
mcp/tools.go (1)
Tool
(69-80)server/session.go (2)
ClientSessionFromContext
(34-39)SessionWithTools
(22-28)
🪛 golangci-lint (1.64.8)
server/session.go
185-185: Error return value of s.SendNotificationToSpecificClient
is not checked
(errcheck)
214-214: Error return value of s.SendNotificationToSpecificClient
is not checked
(errcheck)
server/server.go
809-809: ineffectual assignment to ok
(ineffassign)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 2
♻️ Duplicate comments (4)
server/session.go (4)
68-70
: Potential panic when session ID is unknown
LoadAndDelete
returns(nil, false)
when the key is missing.
The subsequent type‑assertion on anil
interface value (session.(ClientSession)
) will panic.-session, _ := s.sessions.LoadAndDelete(sessionID) -s.hooks.UnregisterSession(ctx, session.(ClientSession)) +sessionVal, ok := s.sessions.LoadAndDelete(sessionID) +if !ok { + return // nothing to do +} +if session, ok := sessionVal.(ClientSession); ok { + s.hooks.UnregisterSession(ctx, session) +}
220-223
: Ignored error fromSendNotificationToSpecificClient
The return value is discarded; if the channel is full or the session is
no longer initialised the error is silently lost, leaving the caller
unaware that the client never received the notification.-// Send notification only to this session -s.SendNotificationToSpecificClient(sessionID, "notifications/tools/list_changed", nil) +// Send notification only to this session +if err := s.SendNotificationToSpecificClient( + sessionID, + "notifications/tools/list_changed", + nil, +); err != nil { + // TODO: decide whether to propagate, retry, or at least log this error +}🧰 Tools
🪛 golangci-lint (1.64.8)
221-221: Error return value of
s.SendNotificationToSpecificClient
is not checked(errcheck)
249-252
: Same unchecked error as above inDeleteSessionTools
Please apply the same handling pattern here to avoid losing delivery
failures.🧰 Tools
🪛 golangci-lint (1.64.8)
250-250: Error return value of
s.SendNotificationToSpecificClient
is not checked(errcheck)
204-218
: 🛠️ Refactor suggestionConcurrent writes to
sessionTools
map can cause data races
AddSessionTools
(and its siblingDeleteSessionTools
) fetch the
map viaGetSessionTools()
and mutate it directly.
BecauseGetSessionTools
returns the underlying map reference, two
goroutines that call these methods for the same session may corrupt the
map concurrently.Consider guarding the per‑session map with a mutex inside the session
implementation or cloning the map before mutation:-sessionTools := session.GetSessionTools() -if sessionTools == nil { - sessionTools = make(map[string]ServerTool) -} - -for _, tool := range tools { - sessionTools[tool.Tool.Name] = tool -} - -session.SetSessionTools(sessionTools) +orig := session.GetSessionTools() +copyMap := maps.Clone(orig) // Go 1.21 +for _, tool := range tools { + copyMap[tool.Tool.Name] = tool +} +session.SetSessionTools(copyMap)(or embed a
sync.Mutex
in the session struct and lock it around
mutations).
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
server/session.go
(1 hunks)server/session_test.go
(1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
server/session.go (4)
mcp/types.go (5)
JSONRPCNotification
(192-195)JSONRPC_VERSION
(81-81)Notification
(106-109)Params
(104-104)NotificationParams
(111-118)server/server.go (2)
ServerTool
(50-53)MCPServer
(126-151)server/errors.go (6)
ErrSessionExists
(16-16)ErrNotificationChannelBlocked
(22-22)ErrNotificationNotInitialized
(21-21)ErrSessionNotFound
(15-15)ErrSessionNotInitialized
(17-17)ErrSessionDoesNotSupportTools
(18-18)mcp/tools.go (1)
Tool
(69-80)
🪛 golangci-lint (1.64.8)
server/session_test.go
442-442: Error return value of server.SendNotificationToSpecificClient
is not checked
(errcheck)
server/session.go
221-221: Error return value of s.SendNotificationToSpecificClient
is not checked
(errcheck)
250-250: Error return value of s.SendNotificationToSpecificClient
is not checked
(errcheck)
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 1
♻️ Duplicate comments (4)
server/session.go (4)
214-221
: Potential data race onsessionTools
mapThe
GetSessionTools
method returns the underlying map, which is then directly modified without synchronization. This could lead to concurrent modifications corrupting the map structure.Consider one of these solutions:
- Protect the per-session tools map with a mutex stored inside the session
- Clone the map before modification using
maps.Clone(sessionTools)
(Go 1.21+)- Serialize tool-modifying operations through the server's event loop
sessionTools := session.GetSessionTools() if sessionTools == nil { sessionTools = make(map[string]ServerTool) +} else { + // Create a copy to avoid data races + newTools := make(map[string]ServerTool, len(sessionTools)) + for k, v := range sessionTools { + newTools[k] = v + } + sessionTools = newTools }
243-250
: Potential data race onsessionTools
map in DeleteSessionToolsSimilar to the issue in
AddSessionTools
, the tools map is directly modified without proper synchronization, which could lead to map corruption under concurrent access.Apply the same map cloning technique as suggested for
AddSessionTools
to avoid race conditions.
226-226
:⚠️ Potential issueIgnored error from
SendNotificationToSpecificClient
hides delivery failuresThe error return value from
SendNotificationToSpecificClient
is not checked, which could hide notification delivery failures.-s.SendNotificationToSpecificClient(sessionID, "notifications/tools/list_changed", nil) +if err := s.SendNotificationToSpecificClient( + sessionID, + "notifications/tools/list_changed", + nil, +); err != nil { + // Consider whether to propagate, retry or log the failure + return fmt.Errorf("failed to send tool list change notification: %w", err) +}🧰 Tools
🪛 golangci-lint (1.64.8)
226-226: Error return value of
s.SendNotificationToSpecificClient
is not checked(errcheck)
255-255
:⚠️ Potential issueUnchecked error from
SendNotificationToSpecificClient
The error return value from
SendNotificationToSpecificClient
is not checked, potentially hiding notification delivery failures.-s.SendNotificationToSpecificClient(sessionID, "notifications/tools/list_changed", nil) +if err := s.SendNotificationToSpecificClient( + sessionID, + "notifications/tools/list_changed", + nil, +); err != nil { + // Consider whether to propagate, retry or log the failure + return fmt.Errorf("failed to send tool list change notification: %w", err) +}🧰 Tools
🪛 golangci-lint (1.64.8)
255-255: Error return value of
s.SendNotificationToSpecificClient
is not checked(errcheck)
🧹 Nitpick comments (1)
server/session.go (1)
92-113
: Consider limiting concurrency when sending notifications to all clientsThe current implementation iterates through all sessions sequentially, which could block the goroutine if there are many sessions.
Consider limiting concurrency when sending notifications to all clients by using a worker pool or semaphore pattern:
s.sessions.Range(func(k, v any) bool { if session, ok := v.(ClientSession); ok && session.Initialized() { - select { - case session.NotificationChannel() <- notification: - // Successfully sent notification - default: - // Channel is blocked, if there's an error hook, use it - if s.hooks != nil && len(s.hooks.OnError) > 0 { - err := ErrNotificationChannelBlocked - go func(sessionID string) { - ctx := context.Background() - // Use the error hook to report the blocked channel - s.hooks.onError(ctx, nil, "notification", map[string]interface{}{ - "method": method, - "sessionID": sessionID, - }, fmt.Errorf("notification channel blocked for session %s: %w", sessionID, err)) - }(session.SessionID()) - } - } + // Use a separate goroutine for each notification to avoid blocking + go func(session ClientSession) { + select { + case session.NotificationChannel() <- notification: + // Successfully sent notification + default: + // Channel is blocked, if there's an error hook, use it + if s.hooks != nil && len(s.hooks.OnError) > 0 { + err := ErrNotificationChannelBlocked + ctx := context.Background() + // Use the error hook to report the blocked channel + s.hooks.onError(ctx, nil, "notification", map[string]interface{}{ + "method": method, + "sessionID": session.SessionID(), + }, fmt.Errorf("notification channel blocked for session %s: %w", session.SessionID(), err)) + } + } + }(session) } return true })This approach would send notifications concurrently rather than sequentially, improving performance when there are many clients.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
server/session.go
(1 hunks)
🧰 Additional context used
🪛 golangci-lint (1.64.8)
server/session.go
226-226: Error return value of s.SendNotificationToSpecificClient
is not checked
(errcheck)
255-255: Error return value of s.SendNotificationToSpecificClient
is not checked
(errcheck)
🔇 Additional comments (1)
server/session.go (1)
1-258
: Overall session implementation is well structuredThe refactoring to move session management into a dedicated file has resulted in a clean, well-organized implementation. The interfaces are well-defined, context management follows Go idioms, and the notification sending logic handles edge cases appropriately.
The session management code is now more maintainable and follows good separation of concerns. The interfaces provide flexibility for different session implementations while maintaining a consistent API.
🧰 Tools
🪛 golangci-lint (1.64.8)
226-226: Error return value of
s.SendNotificationToSpecificClient
is not checked(errcheck)
255-255: Error return value of
s.SendNotificationToSpecificClient
is not checked(errcheck)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 0
♻️ Duplicate comments (5)
server/session_test.go (1)
451-456
:⚠️ Potential issueCheck the error returned when priming the buffer
Static analysis is still flagging this unchecked error.
If the first send fails (e.g. the session becomes un‑initialised), the subsequent test of a blocked channel becomes a false positive.-// Fill the buffer first to ensure it gets blocked -server.SendNotificationToSpecificClient(session.SessionID(), "first-message", nil) +// Fill the buffer and assert it succeeds +require.NoError( + t, + server.SendNotificationToSpecificClient( + session.SessionID(), + "first-message", + nil, + ), +)🧰 Tools
🪛 golangci-lint (1.64.8)
452-452: Error return value of
server.SendNotificationToSpecificClient
is not checked(errcheck)
server/session.go (4)
23-29
: Interface leaks internal state – consider safer accessors
GetSessionTools
returns the mutable map allowing callers to edit it without
validation or concurrency control. Encapsulate mutations behind dedicated
methods (AddTool
,RemoveTool
,ListTools
) to keep invariants intact.This was raised previously and remains unresolved.
214-223
: 🛠️ Refactor suggestionConcurrent writes to
sessionTools
map can cause data races
GetSessionTools
exposes the underlying map.
When two goroutines callAddSessionTools
/DeleteSessionTools
for the same
session they will mutate the same map without synchronisation, leading to
undefined behaviour.Consider cloning the map before modification or protecting the per‑session
tools with a mutex inside the session implementation.- sessionTools := session.GetSessionTools() - if sessionTools == nil { - sessionTools = make(map[string]ServerTool) - } + // Work on a copy to avoid concurrent map writes + original := session.GetSessionTools() + sessionTools := maps.Clone(original) // go1.21 + if sessionTools == nil { + sessionTools = make(map[string]ServerTool) + }(Or embed a
sync.Mutex
inside the session struct and lock around map
accesses.)
226-229
:⚠️ Potential issuePropagate or handle notification errors
The error returned by
SendNotificationToSpecificClient
is ignored.
If the client’s channel is full the mutation succeeds but the user never
learns that the notification was dropped.Either:
- Bubble the error up to the caller, or
- At minimum, log it via the server hooks.
- s.SendNotificationToSpecificClient(sessionID, "notifications/tools/list_changed", nil) + if err := s.SendNotificationToSpecificClient( + sessionID, + "notifications/tools/list_changed", + nil, + ); err != nil { + // Decide whether to rollback, retry, or log – but do not silently ignore. + return err + }🧰 Tools
🪛 golangci-lint (1.64.8)
226-226: Error return value of
s.SendNotificationToSpecificClient
is not checked(errcheck)
248-256
:⚠️ Potential issueSame silent‑error issue when deleting tools
The notification error is again ignored here; please apply the same fix as
above to maintain consistency and reliability.🧰 Tools
🪛 golangci-lint (1.64.8)
255-255: Error return value of
s.SendNotificationToSpecificClient
is not checked(errcheck)
🧹 Nitpick comments (2)
server/session_test.go (2)
278-286
: Usestrings.HasPrefix
to simplify and clarify the filter implementationThe manual prefix check is correct but unnecessarily verbose and error‑prone. Relying on the standard library improves readability and removes the duplicated
len(prefix)
calculation.- var filtered []mcp.Tool - for _, tool := range tools { - if len(tool.Name) >= len(prefix) && tool.Name[:len(prefix)] == prefix { - filtered = append(filtered, tool) - } - } - return filtered + var filtered []mcp.Tool + for _, tool := range tools { + if strings.HasPrefix(tool.Name, prefix) { + filtered = append(filtered, tool) + } + } + return filtered(You will need to add
strings
to the import list at the top of the file.)
340-342
: Avoid magic numbers when asserting tool namesInstead of hard‑coding
6
, reuse the existingprefix
variable to keep the assertion consistent with the filter definition.- assert.True(t, len(tool.Name) >= 6 && tool.Name[:6] == "allow-", - "Tool should start with 'allow-', got: %s", tool.Name) + assert.True(t, strings.HasPrefix(tool.Name, "allow-"), + "Tool should start with 'allow-', got: %s", tool.Name)(Remember to import
strings
once for the entire file.)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
server/session.go
(1 hunks)server/session_test.go
(1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
server/session.go (4)
mcp/types.go (5)
JSONRPCNotification
(192-195)JSONRPC_VERSION
(81-81)Notification
(106-109)Params
(104-104)NotificationParams
(111-118)server/server.go (2)
ServerTool
(50-53)MCPServer
(126-151)server/errors.go (6)
ErrSessionExists
(16-16)ErrNotificationChannelBlocked
(22-22)ErrNotificationNotInitialized
(21-21)ErrSessionNotFound
(15-15)ErrSessionNotInitialized
(17-17)ErrSessionDoesNotSupportTools
(18-18)mcp/tools.go (1)
Tool
(69-80)
🪛 golangci-lint (1.64.8)
server/session.go
226-226: Error return value of s.SendNotificationToSpecificClient
is not checked
(errcheck)
255-255: Error return value of s.SendNotificationToSpecificClient
is not checked
(errcheck)
server/session_test.go
452-452: Error return value of server.SendNotificationToSpecificClient
is not checked
(errcheck)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Really awesome! 🙌
} | ||
|
||
// Send notification to a specific client | ||
err := s.SendNotificationToSpecificClient( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
❤️
Very excited about this! |
I am on a very long flight tomorrow and my goal is to have landed with a working distributed session POC locally with docker and dynamic tool stuff working. I'll do the basics today in case the plane WiFi fails so I don't need to pull anything down, but this is my plan! Apologies it has taken so long! |
Was very quick to get the basics set up with a custom session struct etc. Should only take minutes to have a working tool filtering version built tomorrow, almost done it already in 5 mins of setup! Unless I encounter problems with execution, I am very confident I can make this can do what I need it to. |
OK @ezynda3 this worked really well - but a couple of questions. When am I meant to call Also, how should I think about when to use Tools versus ServerTools - in general I've had the tools travelling around with their handlers and using ServerTool for things, but the tool filter works with []mcp.Tool which means decomposing the servertools I am storing so far. Not sure I'm doing that right, but that's where I ended up. I also ended up not using Anyway, those are my notes so far - but nothing blocking on my end, just thought you'd appreciate seeing where my head was at. The important thing was I was able to make this work well, including with a stateless http implementation where the user just sends extra headers on each request (so not really long lived session). For the dynamic toolsets though, I have to have the stateful session to know what tools the user has available as they hit different servers - so the stateless thing is more of a temporary experiment. Thanks again for getting this built so rapidly! |
The ServerTool is just the underlying container that associates the tool with a handler. I've added a new AddSessionTool func that resembles the AddTool func for consistency. Right now the session stuff is memory only but there is work being done to be able to store sessions anywhere. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 0
♻️ Duplicate comments (4)
server/session_test.go (1)
490-497
:⚠️ Potential issueCheck the error returned by the priming call
Static analysis is still complaining here because the first call fills the buffer but silently discards an error.
If the channel was already full (e.g. in a future refactor with a buffer size 0) the test would turn into a false-positive.-// Fill the buffer first to ensure it gets blocked -server.SendNotificationToSpecificClient(session.SessionID(), "first-message", nil) +require.NoError( + t, + server.SendNotificationToSpecificClient( + session.SessionID(), + "first-message", + nil, + ), +)🧰 Tools
🪛 golangci-lint (1.64.8)
491-491: Error return value of
server.SendNotificationToSpecificClient
is not checked(errcheck)
server/session.go (3)
22-29
:⚠️ Potential issue
SessionWithTools
still exposes the raw map – data-race riskAs noted in the previous review, returning the underlying
map[string]ServerTool
allows concurrent goroutines to mutate it without synchronisation, leading to undefined behaviour.
Please encapsulate access behind accessor/mutator methods or protect the map with a mutex.
219-229
:⚠️ Potential issueConcurrent writes to
sessionTools
without cloning / locking
AddSessionTools
modifies the map fetched fromGetSessionTools()
directly.
If another goroutine reads the same map (e.g. while servicingtools/list
) a race will occur.A minimal fix is to clone the map before writing:
- sessionTools := session.GetSessionTools() - if sessionTools == nil { - sessionTools = make(map[string]ServerTool) - } + original := session.GetSessionTools() + sessionTools := maps.Clone(original) // Go 1.21+
231-233
:⚠️ Potential issueAlways inspect the error from
SendNotificationToSpecificClient
If the channel is blocked, the error is swallowed and the caller is never notified.
Either bubble the error up or log it.- s.SendNotificationToSpecificClient(sessionID, "notifications/tools/list_changed", nil) + if err := s.SendNotificationToSpecificClient( + sessionID, + "notifications/tools/list_changed", + nil, + ); err != nil { + // TODO: decide whether to propagate, retry or log + }🧰 Tools
🪛 golangci-lint (1.64.8)
231-231: Error return value of
s.SendNotificationToSpecificClient
is not checked(errcheck)
🧹 Nitpick comments (5)
README.md (2)
94-98
: Use consistent bullet style to satisfy markdown-lintMarkdown-lint (
MD004
) flags these list items because they use-
while the rest of the document (and the project style-guide) use*
.
Sticking to a single bullet style improves readability and prevents noisy CI failures.- - [Extras](#extras) - - [Session Management](#session-management) - - [Request Hooks](#request-hooks) - - [Tool Handler Middleware](#tool-handler-middleware) + * [Extras](#extras) + * [Session Management](#session-management) + * [Request Hooks](#request-hooks) + * [Tool Handler Middleware](#tool-handler-middleware)🧰 Tools
🪛 markdownlint-cli2 (0.17.2)
94-94: Unordered list style
Expected: asterisk; Actual: dash(MD004, ul-style)
95-95: Unordered list style
Expected: asterisk; Actual: dash(MD004, ul-style)
96-96: Unordered list style
Expected: asterisk; Actual: dash(MD004, ul-style)
97-97: Unordered list style
Expected: asterisk; Actual: dash(MD004, ul-style)
98-98: Unordered list style
Expected: asterisk; Actual: dash(MD004, ul-style)
523-530
: Align list markers with the rest of the READMEThe “Session Management” feature list is also flagged by
markdownlint
for mixing-
with*
.
Consider switching to asterisks for consistency:- - Maintain separate state for each connected client - - Register and track client sessions - - Send notifications to specific clients - - Provide per-session tool customization + * Maintain separate state for each connected client + * Register and track client sessions + * Send notifications to specific clients + * Provide per-session tool customization🧰 Tools
🪛 markdownlint-cli2 (0.17.2)
526-526: Unordered list style
Expected: asterisk; Actual: dash(MD004, ul-style)
527-527: Unordered list style
Expected: asterisk; Actual: dash(MD004, ul-style)
528-528: Unordered list style
Expected: asterisk; Actual: dash(MD004, ul-style)
529-529: Unordered list style
Expected: asterisk; Actual: dash(MD004, ul-style)
server/session_test.go (1)
498-525
: Flaky timing-based assertion – prefer a sync primitiveRelying on
time.Sleep
to wait for the error-hook goroutine can introduce flakiness on slower CI machines.
A simplesync.WaitGroup
or a buffered channel acknowledgement makes the test deterministic.server/session.go (2)
77-115
: Consider removing or flagging unresponsive sessions
SendNotificationToAllClients
merely logs when a channel is blocked, causing the same error every broadcast.
Evicting the session (or downgrading to best-effort delivery) prevents log-spam and keeps the session map healthy.
144-151
: Preserve caller context in async error hookSpawning the goroutine with
context.Background()
loses valuable tracing information.
Pass the originatingctx
instead:-go func(sessionID string) { - ctx := context.Background() +go func(ctx context.Context, sessionID string) { ... -}(session.SessionID()) +}(ctx, session.SessionID())
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
.gitignore
(1 hunks)README.md
(2 hunks)server/session.go
(1 hunks)server/session_test.go
(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- .gitignore
🧰 Additional context used
🪛 markdownlint-cli2 (0.17.2)
README.md
94-94: Unordered list style
Expected: asterisk; Actual: dash
(MD004, ul-style)
95-95: Unordered list style
Expected: asterisk; Actual: dash
(MD004, ul-style)
96-96: Unordered list style
Expected: asterisk; Actual: dash
(MD004, ul-style)
97-97: Unordered list style
Expected: asterisk; Actual: dash
(MD004, ul-style)
526-526: Unordered list style
Expected: asterisk; Actual: dash
(MD004, ul-style)
527-527: Unordered list style
Expected: asterisk; Actual: dash
(MD004, ul-style)
528-528: Unordered list style
Expected: asterisk; Actual: dash
(MD004, ul-style)
529-529: Unordered list style
Expected: asterisk; Actual: dash
(MD004, ul-style)
🪛 golangci-lint (1.64.8)
server/session.go
231-231: Error return value of s.SendNotificationToSpecificClient
is not checked
(errcheck)
260-260: Error return value of s.SendNotificationToSpecificClient
is not checked
(errcheck)
server/session_test.go
491-491: Error return value of server.SendNotificationToSpecificClient
is not checked
(errcheck)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 2
♻️ Duplicate comments (2)
server/session.go (2)
22-29
:SessionWithTools
still exposes raw map → same data-race risk as beforeThe interface keeps returning the underlying
map[string]ServerTool
. Any goroutine that receives this map can mutate it concurrently withAddSessionTools
/DeleteSessionTools
, re-introducing the race that was already flagged in the previous review. Consider exposing accessor methods that avoid handing out the mutable map directly (see earlier suggestion).
238-240
:⚠️ Potential issueIgnored error from
SendNotificationToSpecificClient
The call’s error return is discarded, so a blocked channel or missing session will go unnoticed.
At minimum, log the error; even better, bubble it up to the caller.-// Send notification only to this session -s.SendNotificationToSpecificClient(sessionID, "notifications/tools/list_changed", nil) +// Send notification only to this session +if err := s.SendNotificationToSpecificClient( + sessionID, + "notifications/tools/list_changed", + nil, +); err != nil { + // TODO: decide on retry / propagation strategy + s.logger.Warn("failed to notify session tools change", "sessionID", sessionID, "err", err) +}🧰 Tools
🪛 golangci-lint (1.64.8)
239-239: Error return value of
s.SendNotificationToSpecificClient
is not checked(errcheck)
🧹 Nitpick comments (1)
server/session.go (1)
219-229
: Unnecessary nil-map guard & missed optimisation
range
over anil
map is a no-op, so theif sessionTools != nil
guard can be dropped.
This lets the compiler inline the loop and keeps the code shorter.-// Copy existing tools -if sessionTools != nil { - for k, v := range sessionTools { - newSessionTools[k] = v - } -} +for k, v := range sessionTools { // ok when sessionTools==nil + newSessionTools[k] = v +}🧰 Tools
🪛 golangci-lint (1.64.8)
225-225: S1031: unnecessary nil check around range
(gosimple)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
server/session.go
(1 hunks)
🧰 Additional context used
🪛 golangci-lint (1.64.8)
server/session.go
239-239: Error return value of s.SendNotificationToSpecificClient
is not checked
(errcheck)
277-277: Error return value of s.SendNotificationToSpecificClient
is not checked
(errcheck)
225-225: S1031: unnecessary nil check around range
(gosimple)
🔇 Additional comments (1)
server/session.go (1)
68-75
: Good defensive fix for nil-session panicThe new nil-checks around
LoadAndDelete
remove the panic path reported earlier – looks solid.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 1
♻️ Duplicate comments (6)
server/session.go (6)
26-28
: Consider encapsulating the session tools map to prevent direct manipulation and race conditions.The current interface exposes the internal tools map through getter and setter methods, which allows direct manipulation and could lead to data races with concurrent access.
Consider refactoring the interface to use safer methods that don't expose the underlying map:
// SessionWithTools is an extension of ClientSession that can store session-specific tool data type SessionWithTools interface { ClientSession - // GetSessionTools returns the tools specific to this session, if any - GetSessionTools() map[string]ServerTool - // SetSessionTools sets tools specific to this session - SetSessionTools(tools map[string]ServerTool) + // GetTool returns a specific tool by name, or nil if not found + GetTool(name string) (ServerTool, bool) + // ListTools returns a copy of all tools for this session + ListTools() []ServerTool + // AddTool adds or replaces a tool for this session + AddTool(tool ServerTool) + // RemoveTool removes a tool from this session by name + RemoveTool(name string) bool }This approach encapsulates the storage implementation details and provides safer, more atomic operations.
144-155
: 🛠️ Refactor suggestionFix inconsistent references to error handling hooks.
The code has the same inconsistency as noted earlier, checking
len(s.hooks.OnError)
but callinghooks.onError()
.- hooks.onError(ctx, nil, "notification", map[string]interface{}{ + hooks.OnError[0](ctx, nil, "notification", map[string]interface{}{
198-202
: 🛠️ Refactor suggestionFix inconsistent references to error handling hooks.
The code has the same inconsistency as noted earlier, checking
len(s.hooks.OnError)
but callinghooks.onError()
.- hooks.onError(ctx, nil, "notification", map[string]interface{}{ + hooks.OnError[0](ctx, nil, "notification", map[string]interface{}{
252-256
: 🛠️ Refactor suggestionFix inconsistent references to error handling hooks.
The code has the same inconsistency as noted earlier, checking
len(s.hooks.OnError)
but callinghooks.onError()
.- hooks.onError(ctx, nil, "notification", map[string]interface{}{ + hooks.OnError[0](ctx, nil, "notification", map[string]interface{}{
99-112
: 🛠️ Refactor suggestionFix inconsistent references to error handling hooks.
The code checks
len(s.hooks.OnError)
but callshooks.onError()
. This mismatch in naming could lead to confusion during future refactoring.- if s.hooks != nil && len(s.hooks.OnError) > 0 { + if s.hooks != nil && len(s.hooks.OnError) > 0 { err := ErrNotificationChannelBlocked // Copy hooks pointer to local variable to avoid race condition hooks := s.hooks go func(sessionID string, hooks *Hooks) { ctx := context.Background() // Use the error hook to report the blocked channel - hooks.onError(ctx, nil, "notification", map[string]interface{}{ + hooks.OnError[0](ctx, nil, "notification", map[string]interface{}{ "method": method, "sessionID": sessionID, }, fmt.Errorf("notification channel blocked for session %s: %w", sessionID, err))Make similar changes to all other instances where hooks.onError is used (lines 150, 198, 252, and 303).
303-307
: 🛠️ Refactor suggestionFix inconsistent references to error handling hooks.
The code has the same inconsistency as noted earlier, checking
len(s.hooks.OnError)
but callinghooks.onError()
.- hooks.onError(ctx, nil, "notification", map[string]interface{}{ + hooks.OnError[0](ctx, nil, "notification", map[string]interface{}{
🧹 Nitpick comments (3)
server/session.go (3)
231-235
: Remove unnecessary nil check before ranging over map.A nil check before ranging over a map is redundant since ranging over a nil map in Go is valid and produces no iterations.
// Copy existing tools - if sessionTools != nil { - for k, v := range sessionTools { - newSessionTools[k] = v - } - } + for k, v := range sessionTools { + newSessionTools[k] = v + }🧰 Tools
🪛 golangci-lint (1.64.8)
231-231: S1031: unnecessary nil check around range
(gosimple)
276-278
: Remove unnecessary early return for nil session tools.If the session tools map is nil, continuing execution is safe since ranging over a nil map is valid in Go and creates no iterations.
- if sessionTools == nil { - return nil - }
77-116
: Consider adding a context parameter to SendNotificationToAllClients.All other notification methods accept a context parameter, but this method creates a background context when needed. For consistency and to allow propagating request contexts, consider adding a context parameter.
// SendNotificationToAllClients sends a notification to all the currently active clients. func (s *MCPServer) SendNotificationToAllClients( + ctx context.Context, method string, params map[string]any, ) { // ... // Use the provided context instead of creating a background context - ctx := context.Background() // ... // ... }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
server/session.go
(1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
server/session.go (5)
mcp/types.go (5)
JSONRPCNotification
(206-209)JSONRPC_VERSION
(95-95)Notification
(120-123)Params
(118-118)NotificationParams
(125-132)server/server.go (3)
ServerTool
(50-53)MCPServer
(126-151)ToolHandlerFunc
(41-41)server/errors.go (6)
ErrSessionExists
(16-16)ErrNotificationChannelBlocked
(22-22)ErrNotificationNotInitialized
(21-21)ErrSessionNotFound
(15-15)ErrSessionNotInitialized
(17-17)ErrSessionDoesNotSupportTools
(18-18)server/hooks.go (1)
Hooks
(87-111)mcp/tools.go (1)
Tool
(69-80)
🪛 golangci-lint (1.64.8)
server/session.go
231-231: S1031: unnecessary nil check around range
(gosimple)
Summary by CodeRabbit