Skip to content

Commit d65ddce

Browse files
authored
Fix HandlesIProgressParameter test sporadically failing in CI (#327)
1 parent b0cd654 commit d65ddce

File tree

1 file changed

+14
-3
lines changed

1 file changed

+14
-3
lines changed

tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsToolsTests.cs

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -596,11 +596,22 @@ public async Task HandlesIProgressParameter()
596596

597597
McpClientTool progressTool = tools.First(t => t.Name == nameof(EchoTool.SendsProgressNotifications));
598598

599+
TaskCompletionSource tcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
600+
int remainingNotifications = 10;
601+
599602
ConcurrentQueue<ProgressNotification> notifications = new();
600603
await using (client.RegisterNotificationHandler(NotificationMethods.ProgressNotification, (notification, cancellationToken) =>
601604
{
602-
ProgressNotification pn = JsonSerializer.Deserialize<ProgressNotification>(notification.Params, McpJsonUtilities.DefaultOptions)!;
603-
notifications.Enqueue(pn);
605+
if (JsonSerializer.Deserialize<ProgressNotification>(notification.Params, McpJsonUtilities.DefaultOptions) is { } pn &&
606+
pn.ProgressToken == new ProgressToken("abc123"))
607+
{
608+
notifications.Enqueue(pn);
609+
if (Interlocked.Decrement(ref remainingNotifications) == 0)
610+
{
611+
tcs.SetResult();
612+
}
613+
}
614+
604615
return default;
605616
}))
606617
{
@@ -613,8 +624,8 @@ public async Task HandlesIProgressParameter()
613624
},
614625
cancellationToken: TestContext.Current.CancellationToken);
615626

627+
await tcs.Task;
616628
Assert.Contains("done", JsonSerializer.Serialize(result, McpJsonUtilities.DefaultOptions));
617-
SpinWait.SpinUntil(() => notifications.Count == 10, TimeSpan.FromSeconds(10));
618629
}
619630

620631
ProgressNotification[] array = notifications.OrderBy(n => n.Progress.Progress).ToArray();

0 commit comments

Comments
 (0)