fix: dispatch client request handlers concurrently (#2489)#2490
Open
demoray wants to merge 1 commit intomodelcontextprotocol:mainfrom
Open
fix: dispatch client request handlers concurrently (#2489)#2490demoray wants to merge 1 commit intomodelcontextprotocol:mainfrom
demoray wants to merge 1 commit intomodelcontextprotocol:mainfrom
Conversation
…col#2489) BaseSession._receive_loop awaited each incoming request handler inline, serializing server->client requests (e.g. concurrent sampling calls via asyncio.gather peaked at one in flight). Add an opt-in '_dispatch_requests_concurrently' flag on BaseSession that spawns each request handler in the session's task group. ClientSession enables it; ServerSession stays serial to preserve the initialize ordering that its state machine relies on. Also fix two RequestResponder races that concurrent dispatch widens: - __enter__ no longer replaces the cancel scope, so a cancel() that arrives before the handler enters the context targets the same scope the handler will later run under. - cancel() is idempotent and safe to call before entry. Handler exceptions are translated into a JSON-RPC error response so a raising handler can't wedge the peer.
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
1 participant
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.
Closes #2489.
Problem
BaseSession._receive_loopawaited each incoming request handler inline, so concurrent server→client requests (e.g. Ncreate_messagesampling callbacks launched viaasyncio.gather) ran strictly one-at-a-time. Peak in-flight handlers = 1 regardless of fan-out.Fix
Add an opt-in
_dispatch_requests_concurrentlyflag onBaseSessionthat spawns each request handler in the session's task group viastart_sooninstead of awaiting it inline.ClientSessionenables the flag.ServerSessionstays serial because its_received_requesthas anInitializeRequeststate machine (Initializing → respond → Initialized) that would race if a follow-up request were dispatched concurrently before the state transitioned.Collateral fixes
Concurrent dispatch widens two latent
RequestResponderraces:__enter__used to replace the cancel scope created in__init__. With concurrent dispatch, aCancelledNotificationcan arrive before the handler task enterswith responder:— the pre-entrycancel()would then target a scope that is about to be thrown away. Fixed by not replacing the scope in__enter__.cancel()is now idempotent (early return when already completed) and no longer requires_entered.Handler exceptions are translated into a JSON-RPC error response so a raising handler can't wedge the peer (otherwise the exception would propagate from the dispatch task and tear down the session's task group).
Tests
Four regression tests in
tests/shared/test_session.py:test_concurrent_server_to_client_requests_run_in_parallel— N=4 fan-out, asserts peak concurrency = 4.test_sampling_callback_exception_returns_error_response— handler that raises still gets a JSON-RPC error reply for its id.test_double_cancel_does_not_send_second_response— idempotent cancel.test_cancel_before_context_entered_marks_scope_cancelled— cancel is safe before__enter__.test_handler_that_responds_then_raises_emits_no_duplicate_error— responded-then-raised handlers don't produce a second error for the same id.All 1174 tests pass; 100% branch coverage maintained;
strict-no-cover, ruff, and pyright clean.