Skip to content

Commit 719eade

Browse files
tarekghTarek Mahmoud Sayed
andauthored
Fix flaky AddIncomingMessageFilter_Multiple_Filters_Execute_In_Order test (#1627)
Co-authored-by: Tarek Mahmoud Sayed <tarekms@ntdev.microsoft.com>
1 parent 7ba084a commit 719eade

1 file changed

Lines changed: 15 additions & 0 deletions

File tree

tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsMessageFilterTests.cs

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,12 @@ public async Task AddIncomingMessageFilter_Intercepts_Request_Messages()
100100
[Fact]
101101
public async Task AddIncomingMessageFilter_Multiple_Filters_Execute_In_Order()
102102
{
103+
// The client sends notifications/initialized fire-and-forget, so unlike the initialize and
104+
// tools/list request/response exchanges it has no synchronization point the test can await.
105+
// Signal once the outermost filter finishes processing it so the strict counts below observe a
106+
// complete, stable log instead of racing the still-in-flight notification.
107+
var initializedNotificationProcessed = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
108+
103109
McpServerBuilder
104110
.WithMessageFilters(filters =>
105111
{
@@ -109,6 +115,11 @@ public async Task AddIncomingMessageFilter_Multiple_Filters_Execute_In_Order()
109115
logger.LogInformation("MessageFilter1 before");
110116
await next(context, cancellationToken);
111117
logger.LogInformation("MessageFilter1 after");
118+
119+
if (context.JsonRpcMessage is JsonRpcNotification { Method: NotificationMethods.InitializedNotification })
120+
{
121+
initializedNotificationProcessed.TrySetResult(true);
122+
}
112123
});
113124

114125
filters.AddIncomingFilter((next) => async (context, cancellationToken) =>
@@ -127,6 +138,10 @@ public async Task AddIncomingMessageFilter_Multiple_Filters_Execute_In_Order()
127138

128139
await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
129140

141+
// Wait for the fire-and-forget initialized notification to flow through the filter pipeline
142+
// before snapshotting the log; otherwise the strict counts below can race the notification.
143+
await initializedNotificationProcessed.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken);
144+
130145
var logMessages = MockLoggerProvider.LogMessages
131146
.Where(m => m.Category.StartsWith("MessageFilter"))
132147
.Select(m => m.Message)

0 commit comments

Comments
 (0)