Skip to content
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.

Sign up for GitHub

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

Jump to bottom

fix: handle internal notifications during session cleanup #85

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

fix: handle internal notifications during session cleanup #85

wants to merge 6 commits into from
+34 −10

Conversation

donghao1393
Copy link

@donghao1393 donghao1393 commented Dec 2, 2024 •
edited
Loading

fix: handle internal notifications during session cleanup

Motivation and Context

Addresses an issue where internal notifications (e.g. 'cancelled') during session cleanup would trigger validation errors and crash the server.

How Has This Been Tested?

yes

Breaking Changes

no

Types of changes

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to change)
  • Documentation update

Checklist

  • I have read the MCP Documentation
  • My code follows the repository's style guidelines
  • New and existing tests pass locally
  • I have added appropriate error handling
  • I have added or updated documentation as needed

Additional context

Changes

  • Add session state tracking via _closed flag
  • Skip validation for internal notifications during cleanup phase
  • Improve progress context cleanup with proper async context manager
  • Add final_progress notification to ensure completion state is reported

Testing

Successfully tested with multiple concurrent requests with progress tracking:

  • Cleanup proceeds without validation errors
  • Progress notifications complete properly
  • No server crashes during shutdown

@donghao1393 donghao1393 changed the title fix: improve session cleanup and progress handling fix: handle internal notifications during session cleanup Dec 11, 2024
@rusiaaman rusiaaman mentioned this pull request Dec 15, 2024
Closed
9 tasks
dsp-ant
dsp-ant requested changes Dec 16, 2024
View reviewed changes
Copy link
Member

@dsp-ant dsp-ant left a comment

Choose a reason for hiding this comment

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

Thank you for working on this. I think it's the right direction but needs a bit more refinement. If you don't mind taking a look at my requested changes. I think we are close to getting this in.

src/mcp/shared/session.py Outdated Show resolved Hide resolved
src/mcp/shared/session.py Outdated Show resolved Hide resolved
src/mcp/shared/session.py Outdated Show resolved Hide resolved
@donghao1393
Copy link
Author

Hi @dsp-ant , thank you for your advice. I have made commit. Would you help to check again?

@dsp-ant dsp-ant self-requested a review December 19, 2024 10:12
@donghao1393
Copy link
Author

Technical Details

  • The issue occurred when internal notifications were received during cleanup
  • Added _closed flag to track session state accurately
  • Modified notification handling to be cleanup-aware
  • Ensures graceful shutdown even with pending notifications

Impact

  • Prevents validation errors during shutdown
  • Maintains session stability during cleanup
  • No interference with normal operation
  • Compatible with existing notification patterns

This approach keeps the fix focused on the core issue while maintaining
backward compatibility.

@txbm
Copy link

txbm commented Jan 6, 2025

@dsp-ant @donghao1393 - if I may suggest a more straight-forward fix for this inline with codebase patterns? - txbm@1918562

@donghao1393
Copy link
Author

donghao1393 commented Jan 6, 2025 •
edited
Loading

@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 ?

@orliesaurus
Copy link

Please can you merge this fix in, it's breaking my server on claude desktop app - thank you!

jerome3o-anthropic
jerome3o-anthropic requested changes Jan 18, 2025
View reviewed changes
Copy link
Contributor

@jerome3o-anthropic jerome3o-anthropic left a comment

Choose a reason for hiding this comment

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

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?

src/mcp/shared/session.py Outdated Show resolved Hide resolved
src/mcp/shared/session.py Outdated Show resolved Hide resolved
@donghao1393
Copy link
Author

donghao1393 commented Jan 18, 2025 •
edited
Loading

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?

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:

  1. add this of my mcp service
  2. A client that initiates the task and sends a cancel notification, like:
    • 1st prompt: "create an image using dall-e model to depict a dog standing on the moon watching the earth"
    • 2nd prompt (in the same conversation): "create another image that the dog was running"
    • expected: the conversation should return both 2 images
    • problem: if the cancel notification made the service down in 1st prompt, the 2nd prompt would be failed to send request
  3. inspect the log file mcp-server-openai.log so that we can observe the cancellation behavior

This would help validate both the current behavior and any proposed changes.

@Mehdi-Bl
Copy link

Fixed on my side using this patch + extra glue on MCP 1.2.0.
We need the cancellation notification. Because some time Claude is a bit wild and I noticed if I press cancel and a tool is triggered, it not working.
We need to implement properly cancellation that would stop the tool in a graceful way and don't crash server.

@dsp-ant dsp-ant linked an issue Jan 27, 2025 that may be closed by this pull request
MCP disconnects in asyncio operations #144
Open
@donghao1393 donghao1393 force-pushed the fix/handle-cancelled-notifications branch from 0582169 to 2140ba8 Compare January 29, 2025 11:51
@EduPonz
Copy link

EduPonz commented Jan 29, 2025 •
edited
Loading

@dsp-ant @donghao1393 - if I may suggest a more straight-forward fix for this inline with codebase patterns? - txbm@1918562

I'm going to roll with this until the issue is closed, thanks @txbm!

@donghao1393
fix: improve session cleanup and progress handling
4816bed 
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
fix: correct cancellation notification method name to notifications/c…
e702f9e 
…ancelled

According to MCP spec, the cancellation notification method should
be 'notifications/cancelled' instead of 'cancelled'.
@donghao1393
refactor: remove unnecessary try-except in _should_validate_notification
12692ff 
The try-except block was unnecessary since getattr already handles
the case where the attribute doesn't exist by returning None.
@donghao1393
refactor: improve notification handling
18e029d 
- 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
@donghao1393
fix: add new line in the end of files
2fcc61f
@donghao1393 donghao1393 force-pushed the fix/handle-cancelled-notifications branch from 137c3c1 to 2fcc61f Compare February 5, 2025 17:31
@donghao1393
Copy link
Author

Hi @jerome3o-anthropic , I have updated accordingly and passed my test. Would you help to take a review for the next?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@jerome3o-anthropic jerome3o-anthropic jerome3o-anthropic requested changes

@dsp-ant dsp-ant Awaiting requested review from dsp-ant

Requested changes must be addressed to merge this pull request.

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

MCP disconnects in asyncio operations
7 participants
@donghao1393 @txbm @orliesaurus @Mehdi-Bl @EduPonz @jerome3o-anthropic @dsp-ant