Terminate the client message loop when the transport stream finishes - #275
Open
robertoscipionecom wants to merge 2 commits into
Open
Terminate the client message loop when the transport stream finishes#275robertoscipionecom wants to merge 2 commits into
robertoscipionecom wants to merge 2 commits into
Conversation
The message handling loop treats the end of the transport stream as transient: it exits the `for try await`, calls `receive()` again and starts over. No transport behaves that way. Every one of them exposes a single-use stream that finishes only when the connection is over — StdioTransport on EOF, HTTPClientTransport and InMemoryTransport on disconnect, NetworkTransport on disconnect (its internal reconnect reuses the same continuation without ever finishing it). Once the stream has finished, `receive()` hands back the very same finished stream, so the `for` returns immediately and the `while true` starts over immediately: a busy loop that saturates a core for as long as the client object is alive. Measured in a shipping app: six MCP clients left on dead connections, 599% CPU, with only one server process still running. Break out of the loop instead, the same way the error branch already does. The `resourceTemporarilyUnavailable` retry — the only legitimate reason for this loop to repeat — is untouched.
The existing suite never exercised the branch: `MockTransport.receive()` creates a new stream on every call and replaces the stored continuation without finishing the previous one, so in tests the stream never finishes while the client is still connected. `StreamFinishingTransport` hands out one single-use stream and counts how many times the client asks for it. After the stream finishes, a client whose loop has terminated never asks again, so the count stays at 1. Verified against 0.12.1 without the fix: the test fails with a receive count in the thousands after a 100 ms wait.
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.
Problem
Client.connect(transport:)runs its message handling loop as arepeat { for try await … } while true, which treats the end of the transport streamas transient: it exits the
for, callsreceive()again and starts over.No transport behaves that way. Every transport in this repository exposes a
single-use stream that finishes only when the connection is over:
finish()sitesStdioTransportreadLoopon EOF/read error,disconnect()HTTPClientTransportdisconnect()onlyInMemoryTransportdisconnect(), peer disconnectNetworkTransportdisconnect()/isStopping— its internal reconnect reuses the same continuation without finishing itOnce the stream has finished,
receive()hands back the very same finished stream, sothe
forreturns immediately and thewhile truestarts over immediately. The result isa busy loop that saturates a core for as long as the
Clientobject is alive — and sincethe loop's
Taskholds a strong reference to theClient, that is forever unless someonecalls
disconnect().Measured in a shipping macOS app: six MCP clients left on dead stdio connections,
599% CPU across six spinning threads, with only one server process still running.
A
sampleof the process shows all six threads insideclosure #1 in Client.connect(transport:)→AsyncThrowingStream.Iterator.next().The loop already has the right exit (
break, in the genericcatch), but it is onlyreachable when the stream throws something other than
Errno.resourceTemporarilyUnavailable. A stream that finishes cleanly never gets there.Fix
Break out of the loop when the stream finishes, the same way the error branch already
does. The
resourceTemporarilyUnavailableretry — the only legitimate reason for thisloop to repeat — is untouched. The
repeatis labelled so thebreakinside thedoblock is unambiguous to the reader.
Tests
A regression test is included:
Message loop stops when the transport stream finishes.StreamFinishingTransporthands out one single-use stream and counts how many times theclient asks for it. Once the stream finishes, a client whose loop has terminated never
asks again, so the count stays at 1.
0.12.1without the fix: fails — the receive count is in the thousandsafter a 100 ms wait.
Full suite: 552 tests in 40 suites passing (551 before this PR, all unchanged).
The branch was not covered before, which is why the bug survived:
MockTransport.receive()creates a new stream on every call and replaces the stored continuation without finishing
the previous one, so in tests the stream never finishes while the client is still
connected.