Refactor Client context management to avoid concurrency issue#1054
Merged
jlowin merged 1 commit intoPrefectHQ:mainfrom Jul 6, 2025
Merged
Refactor Client context management to avoid concurrency issue#1054jlowin merged 1 commit intoPrefectHQ:mainfrom
jlowin merged 1 commit intoPrefectHQ:mainfrom
Conversation
jlowin
approved these changes
Jul 6, 2025
Member
There was a problem hiding this comment.
Thanks @hopeful0! This is excellent work. I've been trying to figure out why this is cropping up and appreciate the detective work in #1051 -- it stems from additional calls to list_tools() added in the low-level SDK (modelcontextprotocol/python-sdk#993). Essentially, schema validation now requires each tool call to make additional tool lists to load those schemas, which (in the case of proxy servers) forcibly creates the re-entrant situation that created this race condition.
I like your approach a lot - the only change I want to make is to remove the asserts, which I'd prefer to replace with runtime checks. I can't push to your branch so I'll do that in a subsequent PR. Thanks!
This was referenced Jul 6, 2025
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.

According to #1051, the context manager of the Client may cause a hang in concurrent situations. This PR refactors the methods related to the Client's context manager:
__aenter__to_connect._connectto ensure that the session startup is completed within_context_lock._disconnectto ensure that the session cleanup is completed within_context_lock._ready_eventsettings in_session_runner.These changes might seem somewhat aggressive, but I believe they should make the logic of the Client's context management clearer, thereby avoiding hard-to-detect race conditions. I would like to know your thoughts on these modifications.