diff --git a/tests/ModelContextProtocol.Tests/DiagnosticTests.cs b/tests/ModelContextProtocol.Tests/DiagnosticTests.cs index a8195c78c..a3d9b4e9a 100644 --- a/tests/ModelContextProtocol.Tests/DiagnosticTests.cs +++ b/tests/ModelContextProtocol.Tests/DiagnosticTests.cs @@ -32,6 +32,11 @@ await RunConnected(async (client, server) => var tool = tools.First(t => t.Name == "DoubleValue"); await tool.InvokeAsync(new() { ["amount"] = 42 }, TestContext.Current.CancellationToken); }, clientToServerLog); + + // Wait for server-side activities to be exported. The server processes messages + // via fire-and-forget tasks, so activities may not be immediately available + // after the client operation completes. + await WaitForAsync(() => activities.Count(a => a.Kind == ActivityKind.Server) >= 4); } Assert.NotEmpty(activities); @@ -97,6 +102,9 @@ await RunConnected(async (client, server) => await client.CallToolAsync("Throw", cancellationToken: TestContext.Current.CancellationToken); await Assert.ThrowsAsync(async () => await client.CallToolAsync("does-not-exist", cancellationToken: TestContext.Current.CancellationToken)); }, []); + + // Wait for server-side activities to be exported. + await WaitForAsync(() => activities.Count(a => a.Kind == ActivityKind.Server) >= 4); } Assert.NotEmpty(activities); @@ -164,6 +172,9 @@ await RunConnected(async (client, server) => .First(t => t.Name == "DoubleValue"); await tool.InvokeAsync(new() { ["amount"] = 42 }, TestContext.Current.CancellationToken); }, []); + + // Wait for server-side activities to be exported. + await WaitForAsync(() => activities.Count(a => a.Kind == ActivityKind.Server) >= 3); } // The outer activity should have MCP-specific attributes added to it @@ -215,6 +226,15 @@ private static async Task RunConnected(Func action, await serverTask; } + + private static async Task WaitForAsync(Func condition, int timeoutMs = 10_000) + { + using var cts = new CancellationTokenSource(timeoutMs); + while (!condition()) + { + await Task.Delay(10, cts.Token); + } + } } public class LoggingStream : Stream