feat(core): concurrent inbound dispatch behind an initialization gate#884
Open
devcrocod wants to merge 23 commits into
Open
feat(core): concurrent inbound dispatch behind an initialization gate#884devcrocod wants to merge 23 commits into
devcrocod wants to merge 23 commits into
Conversation
…eardown
Introduce a private Connection object swapped atomically via connectionRef
(CAS), giving connect()/doClose() lifecycle rails: connect() now throws
IllegalStateException if already connected (close() first), and a stale
onClose from a superseded transport is a no-op instead of tearing down the
successor connection. Each connection gets its own CoroutineScope
(SupervisorJob + handlerCoroutineContext + CoroutineName("McpProtocol")),
execution semaphore, and in-flight bookkeeping fields, unused by dispatch
until later tasks enable concurrent handling.
No existing test connected the same Protocol/Client/ServerSession instance
twice without closing first, so no test behavior changes were required.
… bound the response wait by the request timeout Protocol.request() previously left result.await() outside the withTimeout block, so a peer that never responded hung the caller forever; only the timeout path ever sent notifications/cancelled, and plain caller cancellation (job.cancel()) left the response handler registered forever and never notified the peer. - Move result.await() inside withTimeout so response-wait is bounded by the same deadline as the send. - Add a CancellationException catch (after TimeoutCancellationException, since that IS a CancellationException) that runs the same cleanup/notify path as timeout, under NonCancellable so the notification send survives the caller's own cancellation; a failed send is logged/reported via onError but never replaces the original cancellation cause. - Never send notifications/cancelled for `initialize` (the spec forbids cancelling it); local cleanup (handler removal, remembering the id) still runs. - Remember recently locally-cancelled outbound request ids in a bounded FIFO (cap 256) so a late response/progress notification for one is trace-logged instead of hitting onError as "unknown message ID". - Move TestProtocol/RecordingTransport out of ProtocolTest.kt into a shared internal ProtocolTestFakes.kt for reuse by the new test file. Deviations from the pre-existing behavior, both intentional: 1. request() now throws McpException(code=REQUEST_TIMEOUT) on timeout instead of rethrowing the raw TimeoutCancellationException. This matches what RequestOptions.timeout's KDoc has always promised. Updated ClientTest.`should handle request timeout`, which wrapped client.listResources() in an outer withTimeout(1) and asserted TimeoutCancellationException — it now asserts McpException with code REQUEST_TIMEOUT, since any TimeoutCancellationException unwinding through request()'s internal withTimeout is converted, regardless of which scope's deadline actually fired. 2. The timeout branch's notification send is now also wrapped in NonCancellable + try/catch (previously a dead transport on timeout would mask the timeout error) — strictly safer than before.
…out via withTimeoutOrNull The previous commit's catch (TimeoutCancellationException) also captured TimeoutCancellationExceptions raised by an *outer* withTimeout enclosing request(), converting the caller's own timeout into McpException(REQUEST_TIMEOUT) and breaking the caller's contract. Use withTimeoutOrNull for request()'s internal deadline instead: only our own expiry returns null (handled as the REQUEST_TIMEOUT McpException, peer notified); an outer scope's TimeoutCancellationException propagates through withTimeoutOrNull and is handled by the CancellationException branch like any other caller cancellation — cleanup + notifications/cancelled, then the ORIGINAL exception is rethrown unchanged. - Revert the previous commit's edit to ClientTest.`should handle request timeout` — the outer withTimeout(1) there now correctly sees its own TimeoutCancellationException again, so the original assertion stands. - New test: `outer withTimeout around request propagates the original timeout exception and notifies the peer`.
…, with inbound cancellation
Adds an admission cap (maxInFlightHandlers) on top of the execution semaphore (maxConcurrentHandlers): dispatchRequest/dispatchNotification reject/drop overflow inline before launch, so the read loop never suspends on admission. Ping and the cancelled/progress/initialized notifications bypass both tiers so the connection cannot self-block.
…rt before completing the POST In JSON-response mode, handlePostRequest previously relied on inline message dispatch: the handler produced its response (and send() completed the HTTP call) before delivery returned. Once handlers run concurrently, delivery returns immediately, Ktor finalizes the call, and the late call.respond is dropped - the client sees an empty non-success response. handlePostRequest now suspends on a per-POST CompletableDeferred that send() completes once every request id delivered by that POST has a response (JSONRPCResponse or JSONRPCError), then answers with the collected JSON (single object or batch, shapes unchanged). Bodies without requests still get 202 immediately; SSE mode is untouched. Client disconnect or transport close cancels the deferred and retires the POST's ids, so nothing leaks and a late response fails loudly at the send() boundary. Behavior-neutral under inline dispatch: the deferred is already completed by the time the route coroutine awaits it.
…Session Client.connect enables concurrent dispatch right after sending notifications/initialized; ServerSession flips it via the onInitializedNotification router hook, before the user's onInitialized callback runs. Until then inbound messages are handled inline in arrival order, so the initialize exchange stays strictly serial. Test adaptations for the now-asynchronous handler execution: - ClientTest: cancellation test asserts the wire-level notifications/cancelled and the cooperative handler cancellation via deferreds; the timeout test's server handler parks until cancelled instead of racing a real-time delay against the virtual-clock timeout; sendRootsListChanged awaits a deferred. - ClientConnectionLoggingTest: thread-safe lists plus eventually/continually instead of immediate count assertions. - ServerSessionInitializeTest: documents why duplicate-initialize ordering stays deterministic (serial phase). New ProtocolConcurrencyIntegrationTest covers the end-to-end behavior: a server tool handler calling roots/list mid-handling and a client sampling handler calling the server mid-handling no longer deadlock over real serial read loops; cancelling a tool call cascades into its nested roots/list request with no response emitted; and pipelined pre-initialized requests are still answered serially in order, with currentRequestHandlerExtra() available in high-level tool handlers afterwards.
Wraps Protocol.connect()'s trailing transport.start() in a try/catch so a failed start (e.g. connection refused on SSE/StreamableHTTP/WebSocket) rolls back the CAS'd connection instead of permanently wedging the Protocol into "already connected". Also: report send failures for error responses via onError (parity with the other error paths), strengthen the setNotificationHandler KDoc about notifications/cancelled specifically disabling inbound cancellation, switch two tests off runTest's virtual clock onto runBlocking where they exercise real-dispatcher paths, and tighten the PR notes for accuracy (initialize cancellation carve-out, pre-init behavior, a Streamable HTTP JSON-mode known limitation).
In enableJsonResponse mode a request cancelled via notifications/cancelled produces no response, so the POST awaited a response that never came and hung until the client disconnected. Retire cancelled request ids as settled without a response so the exchange completes (202 when nothing remains), and drop late responses for already-retired ids quietly instead of raising an error.
A response or progress message racing request() cancellation on a multi-threaded dispatcher could observe the handler already removed but the id not yet recorded, and report it as an unknown message id. Record the cancelled id first so a racing onResponse/onProgress that sees the handler gone also sees the id and stays quiet.
…andler registration Capture the Connection at entry and re-check it after registering the response handler. If close (or a reconnect) linearizes between the two, doClose already snapshotted the pending handlers without this one, so nothing would complete it — unwind and throw CONNECTION_CLOSED instead of blocking until the request timeout.
…ise or leak In JSON mode a handler outlives its POST route, so a client disconnect can retire a request's mappings before the handler's send() runs. Route request notifications before the streamId lookup, treat a missing streamId for a terminal JSON response as a quiet drop, and reclaim a stray requestToResponseMapping under streamMutex. SSE keeps its "no connection established" diagnostic.
… limits request() now drops its response and progress handlers on every exit path, so a non-cancellation error from the initial transport.send no longer leaks them until close(). connect() rejects a non-positive maxConcurrentHandlers and a maxInFlightHandlers below maxConcurrentHandlers with field-named messages, and the handler-context/concurrency options are documented as read once at connect().
…st mappings When a per-request SSE stream is already gone by the time its response settles, send() now retires the request ids quietly rather than raising "No connection established" and stranding requestToStreamMapping/requestToResponseMapping until close().
Contributor
There was a problem hiding this comment.
Pull request overview
This PR updates the core Protocol dispatch model to prevent head-of-line blocking by allowing inbound request/notification handlers to run concurrently after the MCP initialization handshake (while keeping initialization-phase processing serial). It also adds bounded concurrency controls and improves cancellation propagation, including wiring notifications/cancelled to cooperative handler cancellation.
Changes:
- Introduces an initialization gate that flips inbound dispatch from serial → concurrent, with per-connection bounds (
maxConcurrentHandlers,maxInFlightHandlers) and a configurable handler coroutine context. - Adds cooperative inbound and outbound cancellation behavior (including suppressing responses for peer-cancelled requests and quieting late responses/progress for locally-cancelled outbound requests).
- Updates Streamable HTTP transport JSON-response mode to correctly await asynchronously produced responses and handle cancelled in-flight request IDs.
Reviewed changes
Copilot reviewed 21 out of 21 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpServerTransportTest.kt | Adds regression tests for JSON-response mode with async handler completion and cancellation/retirement behavior. |
| kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpServerTransport.kt | Updates routing/collection logic for JSON-response mode to support concurrent handler completion and cancellation retirement. |
| kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSession.kt | Flips the concurrent-dispatch gate upon receiving notifications/initialized on the server side. |
| kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt | Threads new concurrency options through ServerOptions into ProtocolOptions. |
| kotlin-sdk-server/api/kotlin-sdk-server.api | Updates server module API surface for new constructor params and hook exposure. |
| kotlin-sdk-core/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolConcurrencyJvmTest.kt | Adds real-dispatcher JVM tests to validate genuine parallelism beyond virtual-time schedulers. |
| kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolTestFakes.kt | Introduces shared protocol/transport fakes to support new concurrency and cancellation tests. |
| kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolTest.kt | Expands protocol tests for RequestHandlerExtra context propagation, connection lifecycle, and option validation. |
| kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolOutboundCancellationTest.kt | Adds tests ensuring outbound cancellation sends notifications/cancelled and handles late responses safely. |
| kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolOptionsTest.kt | Verifies default values for newly exposed concurrency knobs. |
| kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolConcurrencyTest.kt | Adds extensive virtual-time tests covering bounded concurrency, admission control, and cancellation suppression. |
| kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/types/serializers.kt | Switches method deserialization to a shared resolver for defined vs custom methods. |
| kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/types/methods.kt | Adds Method.from(...) to resolve known defined methods from wire strings. |
| kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Protocol.kt | Implements the core concurrent dispatch gate, bounded concurrency, in-flight tracking, and cancellation handling. |
| kotlin-sdk-core/api/kotlin-sdk-core.api | Updates core module API dump for new options, context element changes, and new protected hooks. |
| kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/Client.kt | Enables concurrent dispatch after the client completes initialization (post-InitializedNotification). |
| kotlin-sdk-client/api/kotlin-sdk-client.api | Updates client module API dump for new ClientOptions constructor parameters. |
| integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSessionInitializeTest.kt | Documents/locks-in deterministic serial ordering during the initialization phase. |
| integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ClientConnectionLoggingTest.kt | Makes logging notification tests robust under concurrent handler execution by awaiting eventual delivery. |
| integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/ProtocolConcurrencyIntegrationTest.kt | Adds end-to-end deadlock regression tests and cancellation cascade scenarios across client/server. |
| integration-test/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/ClientTest.kt | Adjusts client tests for async handler execution and verifies cancellation reaches the server over the wire. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
- send outbound requests on the captured connection transport - retire batched cancellations on the JSON-response request path - make handler-concurrency caps internal and trim review comments
ChannelTransport starts delivering inbound messages as soon as its read loop launches, which happens before start() marks the transport Operational. Connecting both sides concurrently let the client's initialize reach the server read loop first, so the session answered with a transport-state error instead of an InitializeResult and the handshake failed. createSession() does not await the handshake, so creating the session first is enough to order the two.
e5l
approved these changes
Jul 10, 2026
| _responseHandlers.getAndSet(persistentMapOf()) | ||
| private fun doClose(connection: Connection) { | ||
| // A stale onClose from a previous transport must not tear down the successor connection. | ||
| if (!connectionRef.compareAndSet(connection, null)) { |
Contributor
There was a problem hiding this comment.
should we clean recentlyCancelledRequestIds as well?
Comment on lines
+81
to
+85
| * @property handlerCoroutineContext coroutine context used to run inbound request and notification | ||
| * handlers once initialization has completed. Must contain a real dispatching | ||
| * [kotlin.coroutines.ContinuationInterceptor]. `Dispatchers.Unconfined` and unconfined test | ||
| * dispatchers are unsupported, since handler resumptions would then run on the transport read loop | ||
| * and reintroduce head-of-line blocking. Any [kotlinx.coroutines.Job] in this context is ignored. |
Comment on lines
+90
to
94
| public var handlerCoroutineContext: CoroutineContext = Dispatchers.Default, | ||
| // Internal handler-concurrency bounds, read once at connect and not part of the public surface. | ||
| internal var maxConcurrentHandlers: Int = DEFAULT_MAX_CONCURRENT_HANDLERS, | ||
| internal var maxInFlightHandlers: Int = DEFAULT_MAX_IN_FLIGHT_HANDLERS, | ||
| ) |
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.
Make inbound message handling concurrent after the initialization handshake, so a slow or peer-awaiting handler no longer blocks other messages on the same connection
fixes #176, #572
Motivation and Context
Inbound messages were processed serially on the transport read loop. One slow or peer-awaiting request handler (sampling, elicitation, roots) blocked every later message on that connection, including the responses and
notifications/cancelledthose handlers depend on which could deadlock request/response flows.How Has This Been Tested?
New and old tests
Breaking Changes
maxConcurrentHandlers/maxInFlightHandlers.ProtocolOptions/ClientOptions/ServerOptionsgain new parameters with defaultsRequestHandlerExtrais now a coroutine-context element and is no longer user-constructibleTypes of changes
Checklist