Add support for prompts, notifications and context.Context #1
Conversation
This change introduces three enhancements to the MCP server implementation: 1. Prompts Support - Adds new PromptDefinition type for defining prompt templates - Implements prompts/list and prompts/get RPC methods - Adds rate limiting support for prompt processing - Adds validation for required prompt arguments 2. Notifications System - Adds Notifier interface for sending MCP notifications - Implements connNotifier to handle JSON-RPC notifications - Updates tool and prompt handlers to support custom notification sending 3. Context Support - Adds context.Context to tool and prompt execution functions - Enables proper cancellation and timeout handling - Allows for better request lifecycle management Example usage shows how to: - Define and process prompts with arguments - Send notifications during tool/prompt execution - Use rate limiting for both tools and prompts - Handle various error cases and validation Updates test suite to cover new functionality including: - Prompt listing and processing - Notification sending - Rate limiting behavior - Error handling for missing/invalid arguments Breaking Changes: - Tool execution function signature changes to include context and notifier
There was a problem hiding this comment.
Thanks @tsenart! I left some feedback on the changes, generally looks good.
I hadn't thought through contexts properly yet - presumably we'd like stdio servers to trap SIGTERM and cancel contexts appropriately (could be implemented separately). I think there's some interaction with reading from stdin that might be interesting though.
| }, | ||
| }, | ||
| Process: processPrompt, | ||
| RateLimit: rate.NewLimiter(rate.Every(time.Second), 5), |
There was a problem hiding this comment.
The spec doesn't require rate limiting of prompt gets in the same way that it does for tool calls. However having a rate limit seems reasonable.
| return &t | ||
| } | ||
|
|
||
| // Update the computeSHA256 function to send a single notification |
There was a problem hiding this comment.
This comment looks like a TODO item. It could be removed or changed to document the function.
|
|
||
| responses := allResponses(session.Out.Contents()) | ||
| Expect(responses).To(HaveLen(4)) // First two are initialization, last two are the tool response | ||
| Expect(responses[len(responses)-2]).To(MatchJSON(`{"jsonrpc":"2.0","method":"test/notification","params":{"message":"Processing text"}}`)) |
There was a problem hiding this comment.
Ideally we might be able to make ordered assertions on the responses rather than having to refer to them by index.
gbytes.Say has an implicit cursor that it moves forward as assertions pass but I don't think we can combine that easily with the MatchJSON matcher so this may be the best we can do.
| Expect(responses[len(responses)-1]).To(MatchJSON(`{"jsonrpc":"2.0","id":2,"result":{"messages":[{"role":"assistant","content":{"type":"text","text":"Processed: test input"}}]}}`)) | ||
| }) | ||
|
|
||
| // Missing arguments error |
There was a problem hiding this comment.
The comments for these new test examples don't seem necessary as they match the context descriptions?
| It("responds with a prompt error", func() { | ||
| stdin.WriteString(`{"jsonrpc":"2.0","id":2,"method":"prompts/get","params":{"name":"example","arguments":{"text":""}}}`) | ||
| Eventually(session.Out).Should(gbytes.Say("\n")) | ||
| Expect(lastResponse(session.Out.Contents())).To(MatchJSON(`{"jsonrpc":"2.0","id":2,"result":{"messages":[{"role":"assistant","content":{"type":"text","text":"input text cannot be empty"}}]}}`)) |
There was a problem hiding this comment.
It doesn't seem right to me to include the validation or rate limit errors in the generated prompt that is returned. Seems like these should be actual errors?
Do you think the existing tool rate limit error should be surfaced as a JSON-RPC error rather than a tool execution error?
| return mcp.GetPromptResult{ | ||
| Messages: []mcp.PromptMessage{ | ||
| { | ||
| Role: mcp.RoleAssistant, |
There was a problem hiding this comment.
Would a role of user be more natural for a prompt with a single message?
This change introduces three enhancements to the MCP server implementation:
Example usage shows how to:
Updates test suite to cover new functionality including:
Breaking Changes: