Skip to content

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

Merged
merged 14 commits into from
Apr 29, 2025
Merged

Manage tools on a per session basis #179

merged 14 commits into from
Apr 29, 2025

Conversation

ezynda3
Copy link
Contributor

@ezynda3 ezynda3 commented Apr 20, 2025

Summary by CodeRabbit

  • New Features
    • Added detailed documentation on session management, notifications, and per-session tools in the user guide.
    • Introduced support for session-specific tools, allowing customization and dynamic tool availability per client session.
    • Enabled sending notifications to all clients or specific sessions.
    • Added the ability to filter visible tools based on session context.
  • Bug Fixes
    • Improved error handling for session and notification operations, including blocked notification channels.
  • Tests
    • Added comprehensive tests for session management, tool handling, notifications, and tool filtering.
  • Documentation
    • Expanded the README with practical examples and advanced session management guidance.

Copy link
Contributor

coderabbitai bot commented Apr 20, 2025

Caution

Review failed

The pull request is closed.

Walkthrough

This change refactors session management in the MCP server by moving all session-related interfaces, methods, and logic from server.go into a new server/session.go file. The ClientSession and SessionWithTools interfaces, along with session registration, notification, and per-session tool management, are now encapsulated in this new file. The server's exported methods for handling sessions and notifications are removed from server.go and reimplemented in session.go. Additionally, a new error definitions file server/errors.go is introduced to centralize server error variables. The server now supports tool filtering via a new ToolFilterFunc and associated registration mechanism. The README is updated with detailed documentation on session management and related features, and new tests for session management and tool handling are added.

Changes

File(s) Change Summary
README.md Added a detailed "Extras" section covering session management, per-session tools, notifications, tool filtering, and context propagation, with example code and updated Table of Contents.
server/errors.go New file defining exported error variables for common server, session, and notification error cases.
server/server.go Removed all session management interfaces and methods (including ClientSession and related functions); added ToolFilterFunc and tool filter registration; refactored tool listing and invocation to support per-session tools and tool filtering; replaced context session retrieval with server retrieval.
server/session.go New file implementing ClientSession and SessionWithTools interfaces, session registration/unregistration, notification sending (to all or specific sessions), and per-session tool management, including context utilities.
server/session_test.go New test file with comprehensive tests for session management, per-session tool integration, tool filtering, and notification delivery to sessions.
.gitignore Added .opencode and .claude entries to the ignore list.

Possibly related PRs

  • mark3labs/mcp-go#46: Introduces the initial ClientSession interface and per-session notification channels, forming the foundation for the session management logic that is refactored and extended in this PR.
  • mark3labs/mcp-go#60: Adds Initialize() and Initialized() methods to ClientSession and updates notification logic, which are further reorganized and expanded in the new session.go file in this PR.
  • mark3labs/mcp-go#179: Implements the same session management and per-session tool handling features, including identical interfaces, methods, and documentation additions, indicating direct relation and overlap with this PR.

📜 Recent review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between ccb5cdb and f7ef6bc.

📒 Files selected for processing (2)
  • server/session.go (1 hunks)
  • server/session_test.go (1 hunks)
✨ Finishing Touches
  • 📝 Generate Docstrings

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.

❤️ Share
🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

  • Review comments: Directly reply to a review comment made by CodeRabbit. Example:
    • I pushed a fix in commit <commit_id>, please review it.
    • Generate unit testing code for this file.
    • Open a follow-up GitHub issue for this discussion.
  • Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
    • @coderabbitai generate unit testing code for this file.
    • @coderabbitai modularize this function.
  • PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
    • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
    • @coderabbitai read src/utils.ts and generate unit testing code.
    • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
    • @coderabbitai help me debug CodeRabbit configuration file.

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)

  • @coderabbitai pause to pause the reviews on a PR.
  • @coderabbitai resume to resume the paused reviews.
  • @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
  • @coderabbitai full review to do a full review from scratch and review all the files again.
  • @coderabbitai summary to regenerate the summary of the PR.
  • @coderabbitai generate docstrings to generate docstrings for this PR.
  • @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
  • @coderabbitai resolve resolve all the CodeRabbit review comments.
  • @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
  • @coderabbitai help to get help.

Other keywords and placeholders

  • Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
  • Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
  • Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

  • You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
  • Please see the configuration documentation for more information.
  • If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

  • Visit our Documentation for detailed information on how to use CodeRabbit.
  • Join our Discord Community to get help, request features, and share feedback.
  • Follow us on X/Twitter for updates and announcements.

Copy link
Contributor

@coderabbitai coderabbitai bot left a 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 in DeleteSessionTools

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 sample

The snippet compiles for documentation only, but server.WithToolCapabilities(true) enables list‑changed notifications, not generic “session capabilities”.
Consider adding a clarifying comment or using WithToolCapabilities(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: Use strings.HasPrefix for clearer intent

The 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 swallowed

The non‑blocking send prevents the server from being stalled by a slow client,
which is good. However, the default 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

📥 Commits

Reviewing files that changed from the base of the PR and between 9f39a43 and 08ff414.

📒 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 solid

The sentinel errors are clearly grouped, exported, and ready to be wrapped/checked with errors.Is. No concerns from a correctness or maintainability standpoint.

Copy link
Contributor

@coderabbitai coderabbitai bot left a 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 a nil 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 suggestion

Concurrent write to the underlying map[string]ServerTool – possible data race

AddSessionTools fetches the map via GetSessionTools() and mutates it in‑place without any synchronisation.
If two goroutines call AddSessionTools / DeleteSessionTools concurrently for the same session (or a test mutates the map directly), the map can be corrupted, leading to a panic.

Consider:

  1. Guarding the per‑session map with a mutex inside the session implementation, or
  2. Cloning the map before mutation (maps.Clone in Go 1.21), then replacing it via SetSessionTools.

At minimum, document the expectation that the caller provides external synchronisation.


184-186: ⚠️ Potential issue

Ignored errors from SendNotificationToSpecificClient may hide delivery failures

Both AddSessionTools and DeleteSessionTools drop the error returned by SendNotificationToSpecificClient.
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 issue

Potential deadlock: filters executed while read‑lock is held

handleListTools keeps toolFiltersMu.RLock() while invoking every filter.
If a filter ever registers another filter (calls WithToolFilter) 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 issue

Inner ok shadows outer ok – session tools never win

Inside handleToolCall, the short variable declaration

if sessionWithTools, ok := session.(SessionWithTools); ok {

creates a new ok, leaving the outer ok 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 the ineffassign warning.

🧰 Tools
🪛 golangci-lint (1.64.8)

809-809: ineffectual assignment to ok

(ineffassign)

server/session_test.go (1)

341-360: ⚠️ Potential issue

Tests won’t compile – fakeSession type is still missing

TestMCPServer_SendNotificationToSpecificClient instantiates &fakeSession{…} but no such type is defined in this package (only sessionTestClient and fakeSessionWithTools exist).

Either add the missing stub that satisfies ClientSession or reuse sessionTestClient:

- session1 := &fakeSession{
+ session1 := &sessionTestClient{

Repeat for session2 and session3.
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 runners

The 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

📥 Commits

Reviewing files that changed from the base of the PR and between 08ff414 and b6db184.

📒 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)

Copy link
Contributor

@coderabbitai coderabbitai bot left a 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 a nil 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 from SendNotificationToSpecificClient

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 in DeleteSessionTools

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 suggestion

Concurrent writes to sessionTools map can cause data races

AddSessionTools (and its sibling DeleteSessionTools) fetch the
map via GetSessionTools() and mutate it directly.
Because GetSessionTools 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

📥 Commits

Reviewing files that changed from the base of the PR and between b6db184 and 018a745.

📒 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>
Copy link
Contributor

@coderabbitai coderabbitai bot left a 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 on sessionTools map

The 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:

  1. Protect the per-session tools map with a mutex stored inside the session
  2. Clone the map before modification using maps.Clone(sessionTools) (Go 1.21+)
  3. 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 on sessionTools map in DeleteSessionTools

Similar 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 issue

Ignored error from SendNotificationToSpecificClient hides delivery failures

The 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 issue

Unchecked 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 clients

The 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

📥 Commits

Reviewing files that changed from the base of the PR and between 018a745 and c8f040f.

📒 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 structured

The 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)

Copy link
Contributor

@coderabbitai coderabbitai bot left a 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 issue

Check 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 suggestion

Concurrent writes to sessionTools map can cause data races

GetSessionTools exposes the underlying map.
When two goroutines call AddSessionTools / 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 issue

Propagate 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:

  1. Bubble the error up to the caller, or
  2. 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 issue

Same 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: Use strings.HasPrefix to simplify and clarify the filter implementation

The 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 names

Instead of hard‑coding 6, reuse the existing prefix 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

📥 Commits

Reviewing files that changed from the base of the PR and between c8f040f and 80a7a1e.

📒 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)

Copy link
Contributor

@SamMorrowDrums SamMorrowDrums left a 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(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

❤️

@robert-jackson-glean
Copy link
Contributor

Very excited about this!

@SamMorrowDrums
Copy link
Contributor

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!

@SamMorrowDrums
Copy link
Contributor

I'll do the basics today

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.

@SamMorrowDrums
Copy link
Contributor

OK @ezynda3 this worked really well - but a couple of questions.

When am I meant to call UnregisterSession? I was thinking about it while implementing per user tools. The session only needs to exist as long as the context exists for each request, and the session would be serialised/deserialised from a session store on every request. I realise I don't have to register the session either. Just passing it through on the context is enough.

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 SetSessionTools because I'm mostly working with things like read-only and toolsets as comma separated values, so I'm returning the tools when calling GetSessionTools building the actual tools in that function at the moment.

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!

@ezynda3
Copy link
Contributor Author

ezynda3 commented Apr 29, 2025

OK @ezynda3 this worked really well - but a couple of questions.

When am I meant to call UnregisterSession? I was thinking about it while implementing per user tools. The session only needs to exist as long as the context exists for each request, and the session would be serialised/deserialised from a session store on every request. I realise I don't have to register the session either. Just passing it through on the context is enough.

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 SetSessionTools because I'm mostly working with things like read-only and toolsets as comma separated values, so I'm returning the tools when calling GetSessionTools building the actual tools in that function at the moment.

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.

Copy link
Contributor

@coderabbitai coderabbitai bot left a 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 issue

Check 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 risk

As 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 issue

Concurrent writes to sessionTools without cloning / locking

AddSessionTools modifies the map fetched from GetSessionTools() directly.
If another goroutine reads the same map (e.g. while servicing tools/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 issue

Always 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-lint

Markdown-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 README

The “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 primitive

Relying on time.Sleep to wait for the error-hook goroutine can introduce flakiness on slower CI machines.
A simple sync.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 hook

Spawning the goroutine with context.Background() loses valuable tracing information.
Pass the originating ctx 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

📥 Commits

Reviewing files that changed from the base of the PR and between 0814e56 and 739219c.

📒 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)

Copy link
Contributor

@coderabbitai coderabbitai bot left a 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 before

The interface keeps returning the underlying map[string]ServerTool. Any goroutine that receives this map can mutate it concurrently with AddSessionTools / 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 issue

Ignored 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 a nil map is a no-op, so the if 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

📥 Commits

Reviewing files that changed from the base of the PR and between 739219c and 1d95aa0.

📒 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 panic

The new nil-checks around LoadAndDelete remove the panic path reported earlier – looks solid.

Copy link
Contributor

@coderabbitai coderabbitai bot left a 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 suggestion

Fix inconsistent references to error handling hooks.

The code has the same inconsistency as noted earlier, checking len(s.hooks.OnError) but calling hooks.onError().

-				hooks.onError(ctx, nil, "notification", map[string]interface{}{
+				hooks.OnError[0](ctx, nil, "notification", map[string]interface{}{

198-202: 🛠️ Refactor suggestion

Fix inconsistent references to error handling hooks.

The code has the same inconsistency as noted earlier, checking len(s.hooks.OnError) but calling hooks.onError().

-				hooks.onError(ctx, nil, "notification", map[string]interface{}{
+				hooks.OnError[0](ctx, nil, "notification", map[string]interface{}{

252-256: 🛠️ Refactor suggestion

Fix inconsistent references to error handling hooks.

The code has the same inconsistency as noted earlier, checking len(s.hooks.OnError) but calling hooks.onError().

-				hooks.onError(ctx, nil, "notification", map[string]interface{}{
+				hooks.OnError[0](ctx, nil, "notification", map[string]interface{}{

99-112: 🛠️ Refactor suggestion

Fix inconsistent references to error handling hooks.

The code checks len(s.hooks.OnError) but calls hooks.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 suggestion

Fix inconsistent references to error handling hooks.

The code has the same inconsistency as noted earlier, checking len(s.hooks.OnError) but calling hooks.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

📥 Commits

Reviewing files that changed from the base of the PR and between 1d95aa0 and ccb5cdb.

📒 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)

@ezynda3 ezynda3 merged commit b7ab1cf into main Apr 29, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants