fix: handle internal notifications during session cleanup #85
Conversation
1. Add proper session cleanup handling - Track session state with _closed flag - Handle cancellation and cleanup errors gracefully - Skip notification validation during cleanup 2. Improve progress context - Add final_progress method - Send completion notification in finally block - Handle progress cleanup properly This fix addresses issues with session cleanup causing validation errors and improves progress notification reliability.
donghao1393
changed the title
…ancelled According to MCP spec, the cancellation notification method should be 'notifications/cancelled' instead of 'cancelled'.
The try-except block was unnecessary since getattr already handles the case where the attribute doesn't exist by returning None.
- Always validate all notifications to ensure proper model structure - Handle cancellation notifications separately (no forwarding to stream) - Simplify notification type checking and error handling - Improve code structure and documentation
|
Hi @dsp-ant , thank you for your advice. I have made commit. Would you help to check again? |
- Add CancelledNotification class for handling cancelled requests - Add CancelledNotification to ClientNotification union type - Improve type safety for cancelled notifications
- Add _closed flag to track session state - Skip notification validation during cleanup phase - Handle internal notifications gracefully during shutdown This fixes an issue where internal notifications (e.g. 'cancelled') would trigger validation errors during session cleanup.
|
Technical Details
Impact
This approach keeps the fix focused on the core issue while maintaining |
|
@dsp-ant @donghao1393 - if I may suggest a more straight-forward fix for this inline with codebase patterns? - txbm@1918562 |
I have tested this commit on my local and it's working. I think this is a better solution. It fixed type definition for cancelled notifications to match protocol implementation. This allows proper handling of cancellation messages without requiring changes to the notification processing logic. What do you think @dsp-ant ? |
|
Please can you merge this fix in, it's breaking my server on claude desktop app - thank you! |
There was a problem hiding this comment.
Hey @donghao1393 - It seems to me that this can be simplified a lot now that there is the canceled notification types.
I'm going to be doing a loop on notifications in this sdk soon, and for my understanding (and to write some tests) I'd love it if you could also provide an example that reproduces this issue? would it just be a server with a long running task and the client sending of a cancel notification?
| @@ -204,6 +215,12 @@ async def _send_response( | |||
| ) | |||
| await self._write_stream.send(JSONRPCMessage(jsonrpc_response)) | |||
|
|
|||
| def _is_cancellation_notification(self, message_root: JSONRPCNotification) -> bool: | |||
There was a problem hiding this comment.
This should no longer be needed now that there are the correct types defined
There was a problem hiding this comment.
Thanks for the suggestions. I share your instinct about simplifying the code based on types, but I have some concerns:
-
The cancellation notification we receive is actually using "cancelled" as the method name, not "notifications/cancelled". This is why we need the explicit check in
_is_cancellation_notification. -
Even with the type definitions, the runtime behavior might still need this handling because:
- The notification comes from the transport layer before type validation
- We need to handle both formats for backwards compatibility
- The validation would fail before we can process the notification properly
Would it help if I add some comments explaining this in the code? Or would you prefer to first verify this with the test case I mentioned above?
| await self._received_notification(notification) | ||
| await self._incoming_message_stream_writer.send(notification) | ||
| # check if it is a cancellation notification | ||
| if self._is_cancellation_notification(message.root): |
There was a problem hiding this comment.
Can this be significantly simplified now that we have the types defined?
Yes, you're right. A reproduction case would be a server with a long-running task (like image generation) that typically takes more than a few seconds, and the client sending a cancel notification during the process. I can provide a minimal example:
This would help validate both the current behavior and any proposed changes. |
fix: handle internal notifications during session cleanup
Addresses an issue where internal notifications (e.g. 'cancelled') during session cleanup would trigger validation errors and crash the server.
Changes
_closedflagfinal_progressnotification to ensure completion state is reportedTesting
Successfully tested with multiple concurrent requests with progress tracking: