Skip to content

Fix streamable http sampling #693

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.

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

Merged
merged 3 commits into from
May 12, 2025
Merged

Fix streamable http sampling #693

merged 3 commits into from
May 12, 2025

Conversation

ihrpr
Copy link
Contributor

@ihrpr ihrpr commented May 12, 2025

Thank you @jlowin for raising #680 🙏

The StreamableHTTP client was experiencing a deadlock when the server initiated sampling callbacks during request
processing. The issue occurred due to synchronous request handling in the client's post_writer function.

The Deadlock Scenario:

  1. Client sends a tool call request to the server
  2. Server processes the tool and needs to perform sampling, initiating a callback to the client
  3. Server sends the sampling request to the client as a new JSONRPC request
  4. Client's post_writer receives the sampling request but handles it synchronously
  5. Client blocks waiting for the server's response to the sampling request
  6. Server is blocked waiting for the sampling response before it can complete the original tool call
  7. Deadlock: Both client and server are waiting for each other

The circular dependency prevented the sampling response from ever reaching the server, causing the system to hang
indefinitely.

The Solution

Concurrent request handling in the StreamableHTTP client:

  1. Added TaskGroup parameter to post_writer to enable concurrent task execution
  2. Modified request handling logic to process incoming JSONRPC requests (like sampling callbacks) in separate
    tasks using tg.start_soon(), while responses continue to be handled synchronously (notification need to be handled concurrently as we will have other issues like initialisation notifications )
  3. Fixed SSE event handling to properly detect completion and break from the event loop

The fix ensures that request handling doesn't block the message processing pipeline, allowing the client to:

  • Process server-initiated callbacks in separate tasks
  • Continue listening for new messages in the main task
  • Send callback responses without blocking
  • Receive the server's completion of the original request

Follow ups:

  • as we are working on integration tests - all tests for all the features and transports combinations

@jlowin
Copy link
Contributor

jlowin commented May 12, 2025

Thanks for the quick response @ihrpr! Confirm this fixes both my test case and the original reported bug in the FastMCP repo 🎉

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.

Looks allgood to me - one non-blocking question around breaking/awaiting message handling in the client

@@ -309,6 +310,8 @@ async def _handle_sse_response(
else None
),
)
if is_complete:
break
Copy link
Contributor

Choose a reason for hiding this comment

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

why break here?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

_handle_sse_event returns true if there is a response/error. this means we are done and need to stop the stream

@ihrpr ihrpr merged commit c6fb822 into main May 12, 2025
9 checks passed
@ihrpr ihrpr deleted the ihrpr/fix-shttp-sampling branch May 12, 2025 17:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants