feat(server): persist client info in sessions

[ShawkyZ]

Description

This PR adds the ability to persist client information in MCP sessions.

The main motivation is to enable tracking the origin of requests and which MCP Client triggered them. This makes it possible to use this data later on for analytics.

Type of Change

• New feature (non-breaking change that adds functionality)

Checklist

• My code follows the code style of this project
• I have performed a self-review of my own code
• I have added tests that prove my fix is effective or that my feature works

Additional Information

This PR introduces a new interface SessionWithClientInfo that extends the ClientSession interface to provide the ability to store and retrieve client information.

Usage example

retrievedSession := mcpServer.ClientSessionFromContext(ctx)
sessionWithClientInfo, _ := retrievedSession.(mcpServer.SessionWithClientInfo)
clientName := sessionWithClientInfo.GetClientInfo().Name

Summary by CodeRabbit

• New Features
  • Sessions can now store and retrieve client-specific information during initialization, enhancing session capabilities.
• Tests
  • Added integration tests to verify correct handling and storage of client information in sessions.

[coderabbitai]

Walkthrough

This change introduces the SessionWithClientInfo interface, allowing server sessions to store and retrieve client-specific information. The handleInitialize method is updated to set client info during initialization if the session supports it. Implementations and tests for this interface are added to session types and tested for correct integration and concurrency safety.

Changes

File(s)	Change Summary
server/session.go	Added the SessionWithClientInfo interface, extending session capabilities with GetClientInfo() and SetClientInfo(clientInfo mcp.Implementation) methods.
server/server.go	Updated handleInitialize method to accept the full InitializeRequest and set client info on the session if supported.
server/sse.go	Extended sseSession with a clientInfo field and implemented SessionWithClientInfo methods for atomic client info storage and retrieval. Updated interface assertions.
server/stdio.go	Extended stdioSession with a clientInfo field and implemented SessionWithClientInfo methods for atomic client info storage and retrieval. Updated interface assertions.
server/session_test.go	Added sessionTestClientWithClientInfo implementing SessionWithClientInfo and a new integration test verifying client info storage during session initialization.

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 golangci-lint (1.64.8)

Error: you are using a configuration file for golangci-lint v2 with golangci-lint v1: please use golangci-lint v2
Failed executing command with error: you are using a configuration file for golangci-lint v2 with golangci-lint v1: please use golangci-lint v2

Note

⚡️ AI Code Reviews for VS Code, Cursor, Windsurf

CodeRabbit now has a plugin for VS Code, Cursor and Windsurf. This brings AI code reviews directly in the code editor. Each commit is reviewed immediately, finding bugs before the PR is raised. Seamless context handoff to your AI code agent ensures that you can easily incorporate review feedback.
Learn more here.

---

Note

⚡️ Faster reviews with caching

CodeRabbit now supports caching for code and dependencies, helping speed up reviews. This means quicker feedback, reduced wait times, and a smoother review experience overall. Cached data is encrypted and stored securely. This feature will be automatically enabled for all accounts on May 16th. To opt out, configure Review - Disable Cache at either the organization or repository level. If you prefer to disable all data retention across your organization, simply turn off the Data Retention setting under your Organization Settings.
Enjoy the performance boost—your workflow just got faster.

---

📜 Recent review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
Cache: Disabled due to data retention organization setting
Knowledge Base: Disabled due to data retention organization setting

📥 Commits

Reviewing files that changed from the base of the PR and between 99720bb and 03cb5bf.

📒 Files selected for processing (5)

• server/server.go (2 hunks)
• server/session.go (1 hunks)
• server/session_test.go (4 hunks)
• server/sse.go (3 hunks)
• server/stdio.go (3 hunks)

🧰 Additional context used

🧬 Code Graph Analysis (4)

server/session.go (1)

mcp/types.go (1)

• Implementation (468-471)

server/server.go (2)

mcp/types.go (2)

• InitializeRequest (389-398)
• Params (154-154)

server/session.go (1)

• SessionWithClientInfo (43-49)

server/sse.go (2)

mcp/types.go (1)

• Implementation (468-471)

server/session.go (4)

• ClientSession (11-20)
• SessionWithTools (32-40)
• SessionWithLogging (23-29)
• SessionWithClientInfo (43-49)

server/stdio.go (2)

mcp/types.go (3)

• JSONRPCNotification (310-313)
• Implementation (468-471)
• LoggingLevel (724-724)

server/session.go (3)

• ClientSession (11-20)
• SessionWithLogging (23-29)
• SessionWithClientInfo (43-49)

🔇 Additional comments (15)

server/session.go (1)

42-49: Well-designed interface extension for client info persistence

The new SessionWithClientInfo interface follows the established pattern used by other session extensions like SessionWithLogging and SessionWithTools. The methods are clearly named and the comments provide good documentation for their purpose.

server/server.go (2)

492-496: Updated signature to use the request parameter

Good change to accept the mcp.InitializeRequest parameter that was previously ignored. This allows access to client information in the request.

---

544-549: Well-implemented client info persistence during initialization

The implementation correctly:

1. Checks if the session implements the SessionWithClientInfo interface
2. Only calls SetClientInfo if the session supports it
3. Passes the client info from the request parameters

This defensive programming approach ensures backward compatibility with existing session implementations that don't support client info.

server/sse.go (4)

31-32: Good use of atomic.Value for thread-safe client info storage

Using atomic.Value for the clientInfo field ensures thread-safe operations on this data, which is important in a concurrent environment like an SSE server.

---

91-98: Well-implemented GetClientInfo method

The implementation correctly:

1. Retrieves the value atomically
2. Checks if the value exists
3. Performs type assertion
4. Returns an empty implementation if no client info is stored

---

100-102: Simple and effective SetClientInfo method

The method correctly uses atomic store operation to ensure thread safety when setting client info.

---

104-109: Updated interface assertions for type safety

The interface assertions now include SessionWithClientInfo, which ensures at compile time that sseSession correctly implements all methods of this interface.

server/stdio.go (4)

57-57: Consistent addition of client info field

The client info field uses atomic.Value for thread safety, consistent with the implementation in sseSession.

---

78-85: Correct GetClientInfo implementation

The implementation follows the same pattern as in sseSession, checking if the value exists and returning an empty implementation if none is stored.

---

87-89: Thread-safe SetClientInfo method

The method uses atomic.Store to ensure thread safety when setting client info.

---

104-107: Complete interface assertions

The interface assertions are updated to include SessionWithClientInfo, ensuring type safety at compile time.

server/session_test.go (4)

14-15: Well-structured import for client information.

Adding the mcp package import is necessary for the new SessionWithClientInfo interface implementation.

---

103-138: Great implementation of the SessionWithClientInfo interface.

This test client implementation follows the established patterns in the codebase:

• Uses atomic.Value for thread-safe client info storage
• Provides proper implementation of both ClientSession and SessionWithClientInfo interfaces
• Includes proper null checking and type assertion in GetClientInfo

---

178-180: Good use of compile-time interface compliance check.

Adding the interface validation ensures that sessionTestClientWithClientInfo properly implements the SessionWithClientInfo interface, which is a good practice for catching interface implementation errors at compile time.

---

1085-1128: Comprehensive integration test for the new functionality.

This test thoroughly verifies the client info persistence feature:

• Properly registers a session implementing the new interface
• Tests initialization with client info
• Verifies session retrieval from context
• Confirms interface casting works properly
• Validates that client info is correctly stored and retrieved

The test structure follows the project's testing patterns and provides good coverage of the new functionality.

✨ 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

• X
• Mastodon
• Reddit
• LinkedIn

🪧 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.
  • Explain this complex logic.
  • 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 explain this code block.
  • @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 explain its main purpose.
  • @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.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

• @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.

[ShawkyZ]

Hi @ezynda3 @pottekkat
Apologies for the direct ping.
This feature address an issue that is currently blocking one of the features we are working
Would it be possible to take a look and review it when you have a moment?
All checks have passed, I followed the guidelines and style of the project and I'm happy to address any feedback or make any necessary changes.
I appreciate your time and all the effort that goes into maintaining this project.

Thanks!

[pottekkat]

Thank you 🚀

Since this is a non-breaking change, let's merge this to get the functionality in. Any modifications or improvements to the methods that aren't immediate can be addressed later.

@ezynda3 can you take a quick look?