Skip to content

Commit 63c17f7

Browse files
committed
Fix StdioClientSessionTransport shutdown
StreamClientSessionTransport's ReadMessageAsync method does await CleanupAsync in its finally. CleanupAsync awaits _readTask. _readTask _is_ ReadMessageAsync. So it's waiting for itself to complete, which will then always timeout. We may want to refactor how this works more in the future, but for now, make sure it doesn't wait for itself.
1 parent e39e3ed commit 63c17f7

File tree

1 file changed

+10
-2
lines changed

1 file changed

+10
-2
lines changed

src/ModelContextProtocol/Protocol/Transport/StreamClientSessionTransport.cs

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,16 @@ public StreamClientSessionTransport(
2929
_serverInput = serverInput;
3030
EndpointName = endpointName;
3131

32-
// Start reading messages in the background
32+
// Start reading messages in the background. We use the rarer pattern of new Task + Start
33+
// in order to ensure that the body of the task will always see _readTask initialized.
34+
// It is then able to reliably null it out on completion.
3335
Logger.TransportReadingMessages(endpointName);
34-
_readTask = Task.Run(() => ReadMessagesAsync(_shutdownCts.Token), CancellationToken.None);
36+
var readTask = new Task<Task>(
37+
thisRef => ((StreamClientSessionTransport)thisRef!).ReadMessagesAsync(_shutdownCts.Token),
38+
this,
39+
TaskCreationOptions.DenyChildAttach);
40+
_readTask = readTask.Unwrap();
41+
readTask.Start();
3542

3643
SetConnected(true);
3744
}
@@ -117,6 +124,7 @@ private async Task ReadMessagesAsync(CancellationToken cancellationToken)
117124
}
118125
finally
119126
{
127+
_readTask = null;
120128
await CleanupAsync(cancellationToken).ConfigureAwait(false);
121129
}
122130
}

0 commit comments

Comments
 (0)