feat: Implement per-session prompt functionality #459
Conversation
WalkthroughThis change introduces per-session prompt customization to the session management system. It adds a new interface for session-specific prompts, updates the server to merge and prioritize session prompts over global ones, and implements methods for adding, retrieving, and testing session prompts. Documentation and error handling are also updated accordingly. Changes
Possibly related PRs
Suggested labels
Suggested reviewers
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 📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
✨ Finishing Touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (4)
README.md (1)
550-550: LGTM: Documentation accurately reflects new capability.The addition of "per-session prompt customization" appropriately documents the new feature alongside existing tool customization capabilities.
Consider fixing the markdown list style for consistency:
-- Provide per-session prompt customization +* Provide per-session prompt customizationserver/server.go (1)
849-849: Fix typo in comment.The comment mentions "tools" instead of "prompts".
- // Then override with session-specific tools + // Then override with session-specific promptsserver/session_test.go (2)
486-486: Fix copy-paste error in comment.- // Add session-specific tool using the new helper method + // Add session-specific prompt using the new helper method
642-642: Fix copy-paste errors in comments.Multiple comments incorrectly refer to "tools" instead of "prompts".
- // Add session-specific tools to the uninitialized session + // Add session-specific prompts to the uninitialized session- // Now verify that subsequent tool additions will send notifications + // Now verify that subsequent prompt additions will send notifications- // Verify both tools are available + // Verify both prompts are availableAlso applies to: 669-669, 690-690
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
go.sumis excluded by!**/*.sum
📒 Files selected for processing (8)
README.md(1 hunks)go.mod(1 hunks)server/errors.go(1 hunks)server/server.go(2 hunks)server/session.go(2 hunks)server/session_test.go(6 hunks)server/sse.go(3 hunks)server/sse_test.go(1 hunks)
🧰 Additional context used
🧠 Learnings (6)
📓 Common learnings
Learnt from: lariel-fernandes
PR: mark3labs/mcp-go#428
File: www/docs/pages/servers/prompts.mdx:218-234
Timestamp: 2025-06-20T20:39:51.870Z
Learning: In the mcp-go library, the GetPromptParams.Arguments field is of type map[string]string, not map[string]interface{}, so direct string access without type assertions is safe and correct.
Learnt from: leavez
PR: mark3labs/mcp-go#114
File: client/transport/sse.go:137-179
Timestamp: 2025-04-06T10:07:06.685Z
Learning: The SSE client implementation in the MCP-Go project uses a 30-second timeout for reading SSE events to match the behavior of the original implementation before the transport layer refactoring.
server/sse_test.go (2)
Learnt from: octo
PR: mark3labs/mcp-go#149
File: mcptest/mcptest.go:0-0
Timestamp: 2025-04-21T21:26:32.945Z
Learning: In the mcptest package, prefer returning errors from helper functions rather than calling t.Fatalf() directly, giving callers flexibility in how to handle errors.
Learnt from: lariel-fernandes
PR: mark3labs/mcp-go#428
File: www/docs/pages/servers/prompts.mdx:218-234
Timestamp: 2025-06-20T20:39:51.870Z
Learning: In the mcp-go library, the GetPromptParams.Arguments field is of type map[string]string, not map[string]interface{}, so direct string access without type assertions is safe and correct.
server/sse.go (1)
Learnt from: lariel-fernandes
PR: mark3labs/mcp-go#428
File: www/docs/pages/servers/prompts.mdx:218-234
Timestamp: 2025-06-20T20:39:51.870Z
Learning: In the mcp-go library, the GetPromptParams.Arguments field is of type map[string]string, not map[string]interface{}, so direct string access without type assertions is safe and correct.
server/server.go (1)
Learnt from: lariel-fernandes
PR: mark3labs/mcp-go#428
File: www/docs/pages/servers/prompts.mdx:218-234
Timestamp: 2025-06-20T20:39:51.870Z
Learning: In the mcp-go library, the GetPromptParams.Arguments field is of type map[string]string, not map[string]interface{}, so direct string access without type assertions is safe and correct.
server/session.go (1)
Learnt from: lariel-fernandes
PR: mark3labs/mcp-go#428
File: www/docs/pages/servers/prompts.mdx:218-234
Timestamp: 2025-06-20T20:39:51.870Z
Learning: In the mcp-go library, the GetPromptParams.Arguments field is of type map[string]string, not map[string]interface{}, so direct string access without type assertions is safe and correct.
server/session_test.go (4)
Learnt from: lariel-fernandes
PR: mark3labs/mcp-go#428
File: www/docs/pages/servers/prompts.mdx:218-234
Timestamp: 2025-06-20T20:39:51.870Z
Learning: In the mcp-go library, the GetPromptParams.Arguments field is of type map[string]string, not map[string]interface{}, so direct string access without type assertions is safe and correct.
Learnt from: octo
PR: mark3labs/mcp-go#149
File: mcptest/mcptest.go:0-0
Timestamp: 2025-04-21T21:26:32.945Z
Learning: In the mcptest package, prefer returning errors from helper functions rather than calling t.Fatalf() directly, giving callers flexibility in how to handle errors.
Learnt from: floatingIce91
PR: mark3labs/mcp-go#401
File: server/server.go:1082-1092
Timestamp: 2025-06-23T11:10:42.948Z
Learning: In Go MCP server, ServerTool.Tool field is only used for tool listing and indexing, not for tool execution or middleware. During handleToolCall, only the Handler field is used, so dynamic tools don't need the Tool field populated.
Learnt from: xinwo
PR: mark3labs/mcp-go#35
File: mcp/tools.go:107-137
Timestamp: 2025-03-04T06:59:43.882Z
Learning: Tool responses from the MCP server shouldn't contain RawInputSchema, which is why the UnmarshalJSON method for the Tool struct is implemented to handle only the structured InputSchema format.
🧬 Code Graph Analysis (4)
server/sse_test.go (7)
server/hooks.go (1)
Hooks(94-121)server/session.go (1)
ClientSession(11-20)server/server.go (1)
ServerPrompt(56-59)server/sse.go (1)
NewTestServer(302-308)mcp/prompts.go (5)
Prompt(45-53)GetPromptResult(34-39)PromptMessage(86-89)Role(75-75)RoleUser(78-78)mcp/utils.go (1)
NewGetPromptResult(407-415)mcp/types.go (3)
Content(827-829)TextContent(833-838)TextContent(840-840)
server/sse.go (1)
server/server.go (1)
ServerPrompt(56-59)
server/server.go (2)
server/session.go (2)
ClientSessionFromContext(79-84)SessionWithPrompts(42-48)mcp/prompts.go (1)
Prompt(45-53)
server/session.go (4)
server/server.go (3)
ServerPrompt(56-59)MCPServer(138-163)PromptHandlerFunc(38-38)server/errors.go (2)
ErrSessionNotFound(16-16)ErrSessionDoesNotSupportPrompts(20-20)mcp/prompts.go (1)
Prompt(45-53)server/hooks.go (1)
Hooks(94-121)
🪛 markdownlint-cli2 (0.17.2)
README.md
550-550: Unordered list style
Expected: asterisk; Actual: dash
(MD004, ul-style)
🔇 Additional comments (12)
go.mod (1)
12-12: LGTM: Indirect dependency addition is appropriate.The addition of
golang.org/x/sysas an indirect dependency is expected to support the new session-specific prompt management features that use concurrency primitives.server/errors.go (1)
20-20: LGTM: Error definition follows established patterns.The new error variable
ErrSessionDoesNotSupportPromptsis correctly positioned with other session-related errors and follows the established naming convention, providing consistency with the existingErrSessionDoesNotSupportToolserror.server/sse.go (3)
32-32: LGTM: Consistent field addition for prompt storage.The
prompts sync.Mapfield follows the same pattern as the existingtools sync.Map, providing thread-safe storage for session-specific prompts.
78-87: LGTM: GetSessionPrompts implementation follows established pattern.The method correctly iterates over the sync.Map with proper type assertion and error handling, mirroring the implementation of
GetSessionTools.
100-108: LGTM: SetSessionPrompts implementation is thread-safe and correct.The method properly clears existing prompts before setting new ones, following the same pattern as
SetSessionToolsand ensuring thread-safe operations with sync.Map.server/server.go (1)
894-915: LGTM! Session-specific prompt retrieval is well-implemented.The implementation correctly prioritizes session-specific prompts over global ones and avoids variable shadowing issues.
server/session.go (2)
390-441: Well-implemented session prompt management.The implementation correctly follows the established pattern from
AddSessionTools, including proper error handling, concurrent safety with map copying, and non-blocking notification handling.
443-445: LGTM! Consistent wrapper implementation.The wrapper method follows the same pattern as
AddSessionTool.server/session_test.go (4)
140-140: Good comment correction.Thanks for fixing the comment to accurately describe the type.
175-230: Excellent thread-safe test client implementation.The implementation correctly uses mutex protection and defensive copying to ensure thread safety, following the same robust pattern as
sessionTestClientWithTools.
396-429: Well-structured test for session prompt addition.The test comprehensively verifies prompt addition and notification behavior, following the established pattern from tool tests.
845-915: Comprehensive test for session prompt retrieval.The test thoroughly validates that session-specific prompts correctly override global prompts, with proper type assertion and content validation.
| type SessionWithPrompts interface { | ||
| ClientSession | ||
| // GetPrompts returns the prompts specific to this session, if any | ||
| GetSessionPrompts() map[string]ServerPrompt | ||
| // SetPrompts sets prompts specific to this session | ||
| SetSessionPrompts(prompts map[string]ServerPrompt) | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Document thread-safety requirements for consistency.
The SessionWithTools interface above documents thread-safety requirements in its method comments, but SessionWithPrompts does not. For consistency and to ensure proper implementation, please add similar documentation.
type SessionWithPrompts interface {
ClientSession
- // GetPrompts returns the prompts specific to this session, if any
+ // GetSessionPrompts returns the prompts specific to this session, if any
+ // This method must be thread-safe for concurrent access
GetSessionPrompts() map[string]ServerPrompt
- // SetPrompts sets prompts specific to this session
+ // SetSessionPrompts sets prompts specific to this session
+ // This method must be thread-safe for concurrent access
SetSessionPrompts(prompts map[string]ServerPrompt)
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| type SessionWithPrompts interface { | |
| ClientSession | |
| // GetPrompts returns the prompts specific to this session, if any | |
| GetSessionPrompts() map[string]ServerPrompt | |
| // SetPrompts sets prompts specific to this session | |
| SetSessionPrompts(prompts map[string]ServerPrompt) | |
| } | |
| type SessionWithPrompts interface { | |
| ClientSession | |
| // GetSessionPrompts returns the prompts specific to this session, if any | |
| // This method must be thread-safe for concurrent access | |
| GetSessionPrompts() map[string]ServerPrompt | |
| // SetSessionPrompts sets prompts specific to this session | |
| // This method must be thread-safe for concurrent access | |
| SetSessionPrompts(prompts map[string]ServerPrompt) | |
| } |
🤖 Prompt for AI Agents
In server/session.go around lines 42 to 48, the SessionWithPrompts interface
lacks thread-safety documentation unlike the SessionWithTools interface. Add
comments to the GetSessionPrompts and SetSessionPrompts methods specifying their
thread-safety requirements, indicating whether they must be safe for concurrent
use or require external synchronization, to maintain consistency and guide
proper implementation.
Description
This PR adds support for session-specific prompts in the MCP server, allowing different prompts to be available per client session. This mirrors the existing session-specific tools functionality and enables dynamic prompt customization based on session context.
Type of Change
Changes Made
Core Functionality
sseSessionstruct: Addedprompts sync.Mapfield to store session-specific promptsSessionWithPromptsinterface: Added methods for managing session promptshandleGetPrompt: Modified to check session-specific prompts first, then fall back to global promptshandleListPrompts: Modified to merge session-specific prompts with global prompts in the listServer Methods
AddSessionPrompts(sessionID string, prompts ...ServerPrompt) errorAddSessionPrompt(sessionID string, prompt mcp.Prompt, handler PromptHandlerFunc) errorDeleteSessionPrompts(sessionID string, names ...string) errorTesting
Checklist
MCP Spec Compliance
Additional Information
This implementation follows the same pattern as the existing session-specific tools functionality, ensuring consistency across the codebase. The feature allows for:
The implementation maintains full backward compatibility and doesn't affect existing global prompt functionality.
Summary by CodeRabbit
New Features
Documentation
Tests